Land the Discord gateway bot (port + hardening) (#7)

notifications/discord_bot.py

@@ -0,0 +1,295 @@
+"""Discord **gateway** bot — reacts to @mentions (and DMs) in real time.
+
+The other two Discord paths need no persistent connection: slash commands arrive as
+signed HTTP POSTs (discord_interactions.py) and the daily post goes out over an
+incoming webhook (discord.py). 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), so — like the subscription notifier and the daily post — this
+runs in exactly one process: web/app.py starts it only when the process wins the
+leader election (see core/singleton.claim_leader). No new container or daemon; it
+rides the app's asyncio event loop as a background task.
+
+Least privilege / no privileged intent. We subscribe to 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 works with no portal review. To trigger on
+arbitrary keywords in channels, enable the MESSAGE_CONTENT intent in the Developer
+Portal and extend `_response_for_message` — the rest of the machinery is unchanged.
+
+Enable with THERMOGRAPH_DISCORD_BOT=1 (plus THERMOGRAPH_DISCORD_BOT_TOKEN, the same
+bot token discord.py uses for DMs). Unset => no gateway connection; everything here
+no-ops, so an unconfigured deploy pays nothing.
+
+Behaviour: mention the bot with a city name ("@Thermograph Phoenix") and it replies
+with the same grade embed the /grade slash command returns, reusing
+discord_interactions._grade_message so the two never drift.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+import random
+
+from notifications import discord
+
+log = logging.getLogger("thermograph.discord_bot")
+
+# 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.
+_GATEWAY_URL = "wss://gateway.discord.gg/?v=10&encoding=json"
+_GATEWAY_QS = "/?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.
+_INTENT_GUILDS = 1 << 0
+_INTENT_GUILD_MESSAGES = 1 << 9
+_INTENT_DIRECT_MESSAGES = 1 << 12
+INTENTS = _INTENT_GUILDS | _INTENT_GUILD_MESSAGES | _INTENT_DIRECT_MESSAGES
+
+# Gateway opcodes (discord.com/developers/docs/topics/gateway-events#opcodes).
+_OP_DISPATCH = 0
+_OP_HEARTBEAT = 1
+_OP_IDENTIFY = 2
+_OP_RESUME = 6
+_OP_RECONNECT = 7
+_OP_INVALID_SESSION = 9
+_OP_HELLO = 10
+_OP_HEARTBEAT_ACK = 11
+
+# Close codes for unrecoverable conditions — 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.
+_FATAL_CLOSE = {4004, 4010, 4011, 4012, 4013, 4014}
+
+_TRUTHY = {"1", "true", "yes", "on"}
+
+_HELP = ("Mention me with a city name — e.g. `@Thermograph Phoenix` — and I'll grade "
+         "today's weather against ~45 years of local history.")
+
+# Reconnect backoff bounds (seconds).
+_BACKOFF_START = 1.0
+_BACKOFF_MAX = 60.0
+
+
+def enabled() -> bool:
+    """True iff the gateway bot is switched on AND a bot token is configured. Off by
+    default: a gateway connection is a new persistent behaviour, so it takes an
+    explicit opt-in rather than piggybacking on the token being present for DMs."""
+    flag = os.environ.get("THERMOGRAPH_DISCORD_BOT", "").strip().lower() in _TRUTHY
+    return flag and bool(discord.BOT_TOKEN)
+
+
+# --- message handling (pure; unit-tested without a live gateway) --------------
+def _mentions_bot(msg: dict, bot_user_id: str) -> bool:
+    return any(str(u.get("id")) == str(bot_user_id) for u in (msg.get("mentions") or []))
+
+
+def _is_dm(msg: dict) -> bool:
+    # A DM has no guild_id; a guild message always carries one.
+    return not msg.get("guild_id")
+
+
+def _strip_mentions(content: str, bot_user_id: str) -> str:
+    """Drop every form of the bot's own mention (<@id> and the legacy <@!id>) and
+    collapse the surrounding whitespace, leaving just the user's query text."""
+    out = content.replace(f"<@{bot_user_id}>", " ").replace(f"<@!{bot_user_id}>", " ")
+    return " ".join(out.split())
+
+
+def _response_for_message(msg: dict, bot_user_id: str) -> dict | None:
+    """The reply payload for one MESSAGE_CREATE, or None to stay silent.
+
+    Silent unless the message DMs the bot or @mentions it, and never for a bot
+    author (that includes ourselves — the guard against a mention-loop). The text
+    after the mention is treated as a city query and answered with the same embed
+    the /grade slash command builds; an empty query gets the help line."""
+    author = msg.get("author") or {}
+    if author.get("bot") or str(author.get("id")) == str(bot_user_id):
+        return None
+    dm = _is_dm(msg)
+    if not dm and not _mentions_bot(msg, bot_user_id):
+        return None
+    content = msg.get("content") or ""
+    query = content.strip() if dm else _strip_mentions(content, bot_user_id)
+    if not query:
+        return {"content": _HELP}
+    # Import here (not at module top) to avoid importing the FastAPI web layer's
+    # transitive deps for a module that also loads in the notifier-only context.
+    from notifications import discord_interactions
+    data = discord_interactions._grade_message(query)
+    # Gateway messages can't be ephemeral — the flag is interactions-only.
+    data.pop("flags", None)
+    return data
+
+
+# --- REST reply ---------------------------------------------------------------
+async def _send_reply(channel_id, message_id, data: dict) -> None:
+    """Post the reply in-channel as a proper reply to the triggering message.
+
+    Reuses discord._bot_post (sync httpx with a 429 retry) off the event loop via a
+    thread so the gateway read-loop never blocks. allowed_mentions is locked down so
+    a crafted query can't turn the reply into an @everyone/role ping; only the person
+    who asked is pinged, via the reply reference."""
+    if not channel_id:
+        return
+    payload = dict(data)
+    if message_id:
+        payload["message_reference"] = {"message_id": str(message_id)}
+        payload["allowed_mentions"] = {"parse": [], "replied_user": True}
+    try:
+        await asyncio.to_thread(
+            discord._bot_post, f"/channels/{channel_id}/messages", payload)
+    except Exception:  # noqa: BLE001 - a failed reply must never kill the read loop
+        log.warning("discord_bot: reply failed", exc_info=True)
+
+
+# --- gateway session ----------------------------------------------------------
+class _Session:
+    """Mutable per-connection state that survives a resume: the last sequence
+    number, the resume token/URL, and our own user id (for mention detection)."""
+
+    __slots__ = ("seq", "session_id", "resume_url", "user_id")
+
+    def __init__(self) -> None:
+        self.seq: int | None = None
+        self.session_id: str | None = None
+        self.resume_url: str | None = None
+        self.user_id: str = ""
+
+
+async def _heartbeat(ws, interval: float, sess: _Session, acked: dict) -> None:
+    """Send op-1 heartbeats every `interval` seconds (first beat jittered per
+    Discord's guidance). If the previous beat was never ACKed the link is a zombie —
+    close with a non-1000 code so the next connect can resume, and stop."""
+    await asyncio.sleep(interval * random.random())
+    while True:
+        if not acked["ok"]:
+            await ws.close(code=4000)
+            return
+        acked["ok"] = False
+        try:
+            await ws.send(json.dumps({"op": _OP_HEARTBEAT, "d": sess.seq}))
+        except Exception:  # noqa: BLE001 - closed socket; the read loop handles reconnect
+            return
+        await asyncio.sleep(interval)
+
+
+async def _dispatch(msg: dict, sess: _Session) -> None:
+    """Handle an op-0 DISPATCH event we care about."""
+    t = msg.get("t")
+    d = msg.get("d") or {}
+    if t == "READY":
+        sess.session_id = d.get("session_id")
+        resume = (d.get("resume_gateway_url") or "").rstrip("/")
+        sess.resume_url = f"{resume}{_GATEWAY_QS}" if resume else None
+        sess.user_id = str((d.get("user") or {}).get("id") or "")
+        log.info("discord_bot: gateway ready (user %s)", sess.user_id)
+    elif t == "MESSAGE_CREATE" and sess.user_id:
+        # _response_for_message reaches into discord_interactions._grade_message for
+        # anything but the help line -- sync parquet/DB reads + grading math, same as
+        # web/app.py's discord_interactions route -- keep it off the (single, shared)
+        # gateway event loop so one grading lookup can't stall heartbeats/reads.
+        reply = await asyncio.to_thread(_response_for_message, d, sess.user_id)
+        if reply:
+            await _send_reply(d.get("channel_id"), d.get("id"), reply)
+
+
+async def _connection(ws, sess: _Session, resume: bool) -> str:
+    """Drive one live connection until it closes. Returns the mode for the next
+    connect: "resume" to reconnect and RESUME, "identify" to start a fresh session."""
+    hello = json.loads(await ws.recv())
+    if hello.get("op") != _OP_HELLO:
+        return "identify"  # protocol out of step — start clean
+    interval = float(hello["d"]["heartbeat_interval"]) / 1000.0
+    acked = {"ok": True}
+    hb = asyncio.create_task(_heartbeat(ws, interval, sess, acked))
+    try:
+        if resume and sess.session_id and sess.seq is not None:
+            await ws.send(json.dumps({"op": _OP_RESUME, "d": {
+                "token": discord.BOT_TOKEN,
+                "session_id": sess.session_id,
+                "seq": sess.seq,
+            }}))
+        else:
+            await ws.send(json.dumps({"op": _OP_IDENTIFY, "d": {
+                "token": discord.BOT_TOKEN,
+                "intents": INTENTS,
+                "properties": {"os": "linux", "browser": "thermograph", "device": "thermograph"},
+            }}))
+        async for raw in ws:
+            msg = json.loads(raw)
+            if msg.get("s") is not None:
+                sess.seq = msg["s"]
+            op = msg.get("op")
+            if op == _OP_DISPATCH:
+                await _dispatch(msg, sess)
+            elif op == _OP_HEARTBEAT:            # server asked for an immediate beat
+                await ws.send(json.dumps({"op": _OP_HEARTBEAT, "d": sess.seq}))
+            elif op == _OP_HEARTBEAT_ACK:
+                acked["ok"] = True
+            elif op == _OP_RECONNECT:            # planned reconnect — resume
+                return "resume"
+            elif op == _OP_INVALID_SESSION:
+                # d == True => the session can still be resumed; else it's dead and
+                # the next connect must IDENTIFY fresh.
+                if not msg.get("d"):
+                    sess.session_id = None
+                    sess.resume_url = None
+                await asyncio.sleep(1 + random.random())
+                return "resume" if sess.session_id else "identify"
+    finally:
+        hb.cancel()
+    return "resume" if sess.session_id else "identify"
+
+
+def _close_code(exc: Exception):
+    """Best-effort websocket close code across websockets library versions (older
+    exposes .code; newer nests it under .rcvd)."""
+    code = getattr(exc, "code", None)
+    if code:
+        return code
+    return getattr(getattr(exc, "rcvd", None), "code", None)
+
+
+async def run() -> None:
+    """Maintain the gateway connection for the app's lifetime: connect, IDENTIFY (or
+    RESUME), dispatch events, and reconnect with capped exponential backoff. Returns
+    only on a fatal close (bad token / disallowed intents) or task cancellation, and
+    never raises out — a misbehaving gateway must not take the app down."""
+    if not enabled():
+        return
+    try:
+        import websockets
+    except Exception:  # noqa: BLE001 - shipped via uvicorn[standard]; guard anyway
+        log.warning("discord_bot: 'websockets' unavailable; gateway bot disabled")
+        return
+    sess = _Session()
+    mode = "identify"
+    backoff = _BACKOFF_START
+    while True:
+        use_resume = mode == "resume" and bool(sess.resume_url)
+        url = sess.resume_url if use_resume else _GATEWAY_URL
+        try:
+            async with websockets.connect(url, max_size=2 ** 20, open_timeout=30) as ws:
+                mode = await _connection(ws, sess, resume=use_resume)
+            backoff = _BACKOFF_START   # server-requested reconnect: come back promptly
+            continue
+        except asyncio.CancelledError:
+            log.info("discord_bot: shutting down")
+            raise
+        except Exception as exc:  # noqa: BLE001 - any drop -> back off and reconnect
+            if _close_code(exc) in _FATAL_CLOSE:
+                log.error("discord_bot: fatal gateway close %s; stopping", _close_code(exc))
+                return
+            mode = "resume" if sess.session_id else "identify"
+            log.warning("discord_bot: gateway dropped (%s); reconnecting in %.0fs",
+                        type(exc).__name__, backoff)
+        await asyncio.sleep(backoff + random.random())
+        backoff = min(backoff * 2, _BACKOFF_MAX)

requirements.txt

@@ -19,6 +19,9 @@ alembic==1.14.0
 pywebpush==2.0.0
 # Ed25519 verification of Discord interaction webhooks (see notifications/discord_interactions.py).
 PyNaCl==1.5.0
+# Discord gateway bot's websocket client (notifications/discord_bot.py). Already
+# pulled in transitively by uvicorn[standard]; pinned here since we import it directly.
+websockets==16.0
 # Recurring worker-tier jobs — city warming, IndexNow pings (see
 # notifications/scheduler.py). In-process; no broker/queue needed at this scale.
 apscheduler==3.11.3

tests/notifications/test_discord_bot.py

@@ -0,0 +1,125 @@
+"""Discord gateway bot: the message-handling logic that decides when to reply and
+with what. No live gateway — we feed MESSAGE_CREATE payloads straight to the pure
+handler and stub the shared /grade lookup so these stay data-independent."""
+import asyncio
+
+import pytest
+
+from notifications import discord
+from notifications import discord_bot as bot
+from notifications import discord_interactions as di
+
+BOT_ID = "999"
+
+
+@pytest.fixture(autouse=True)
+def _stub_grade(monkeypatch):
+    """Reply payload for /grade, mirroring the real shape (embeds + an ephemeral
+    flag the gateway path must drop)."""
+    monkeypatch.setattr(
+        di, "_grade_message",
+        lambda query: {"embeds": [{"title": f"grade:{query}"}], "flags": 64})
+
+
+def _msg(content="", *, author_id="1", bot_author=False, mentions=(), guild=True):
+    m = {
+        "content": content,
+        "id": "42",
+        "channel_id": "77",
+        "author": {"id": author_id, "bot": bot_author},
+        "mentions": [{"id": mid} for mid in mentions],
+    }
+    if guild:
+        m["guild_id"] = "5"
+    return m
+
+
+# --- when the bot stays silent ----------------------------------------------
+
+def test_ignores_other_bots():
+    assert bot._response_for_message(
+        _msg("<@999> Phoenix", author_id="8", bot_author=True, mentions=[BOT_ID]), BOT_ID) is None
+
+
+def test_ignores_its_own_messages():
+    assert bot._response_for_message(
+        _msg("hello", author_id=BOT_ID, mentions=[BOT_ID]), BOT_ID) is None
+
+
+def test_ignores_guild_message_without_a_mention():
+    assert bot._response_for_message(_msg("just chatting", mentions=[]), BOT_ID) is None
+
+
+# --- when the bot replies ----------------------------------------------------
+
+def test_mention_with_city_returns_grade_without_ephemeral_flag():
+    out = bot._response_for_message(_msg("<@999> Phoenix", mentions=[BOT_ID]), BOT_ID)
+    assert out == {"embeds": [{"title": "grade:Phoenix"}]}   # flags dropped
+    assert "flags" not in out
+
+
+def test_legacy_nickname_mention_is_stripped():
+    out = bot._response_for_message(_msg("<@!999>   Tokyo", mentions=[BOT_ID]), BOT_ID)
+    assert out["embeds"][0]["title"] == "grade:Tokyo"
+
+
+def test_dm_needs_no_mention():
+    out = bot._response_for_message(_msg("Berlin", guild=False, mentions=[]), BOT_ID)
+    assert out["embeds"][0]["title"] == "grade:Berlin"
+
+
+def test_mention_with_no_query_gives_help():
+    out = bot._response_for_message(_msg("<@999>", mentions=[BOT_ID]), BOT_ID)
+    assert "content" in out and "mention me" in out["content"].lower()
+
+
+# --- helpers -----------------------------------------------------------------
+
+def test_strip_mentions_handles_both_forms():
+    assert bot._strip_mentions("<@999> <@!999>  Sao Paulo", BOT_ID) == "Sao Paulo"
+
+
+def test_is_dm_is_true_without_guild_id():
+    assert bot._is_dm(_msg("x", guild=False)) is True
+    assert bot._is_dm(_msg("x", guild=True)) is False
+
+
+# --- _dispatch: same reply, off the event loop -------------------------------
+
+def test_dispatch_delivers_the_same_reply_via_a_thread(monkeypatch):
+    """_dispatch now runs _response_for_message through asyncio.to_thread (the
+    event-loop fix) -- confirm that still produces the exact same reply and gets
+    it to _send_reply, i.e. the offload didn't change behaviour."""
+    sent = {}
+
+    async def _fake_send_reply(channel_id, message_id, data):
+        sent["channel_id"] = channel_id
+        sent["message_id"] = message_id
+        sent["data"] = data
+
+    monkeypatch.setattr(bot, "_send_reply", _fake_send_reply)
+    sess = bot._Session()
+    sess.user_id = BOT_ID
+    msg = {"t": "MESSAGE_CREATE", "d": _msg("<@999> Phoenix", mentions=[BOT_ID])}
+
+    asyncio.run(bot._dispatch(msg, sess))
+
+    assert sent["channel_id"] == "77"
+    assert sent["message_id"] == "42"
+    assert sent["data"] == {"embeds": [{"title": "grade:Phoenix"}]}
+
+
+# --- enabled() gating --------------------------------------------------------
+
+def test_enabled_is_off_by_default(monkeypatch):
+    monkeypatch.delenv("THERMOGRAPH_DISCORD_BOT", raising=False)
+    monkeypatch.setattr(discord, "BOT_TOKEN", "a-token")
+    assert bot.enabled() is False
+
+
+def test_enabled_needs_both_flag_and_token(monkeypatch):
+    monkeypatch.setenv("THERMOGRAPH_DISCORD_BOT", "1")
+    monkeypatch.setattr(discord, "BOT_TOKEN", "")
+    assert bot.enabled() is False
+    monkeypatch.setattr(discord, "BOT_TOKEN", "a-token")
+    assert bot.enabled() is True

web/app.py

@@ -1,4 +1,5 @@
 """Thermograph API — grade recent local weather against ~45 years of climatology."""
+import asyncio
 import contextlib
 import datetime
 import functools
@@ -22,6 +23,7 @@ from api import content_routes
 from core import audit
 from data import climate
 from notifications import digest
+from notifications import discord_bot
 from notifications import discord_interactions as discord_interactions_mod
 from notifications import discord_link
 from accounts import db
@@ -162,16 +164,31 @@ async def _lifespan(app):
     # restricts this to processes that are allowed to own it at all — "web" replicas
     # never start it even if they'd otherwise win the leader election, so a stateless
     # N-replica web tier can scale without also scaling notifier instances.
+    bot_task = None
     if _should_run_notifier():
         notify.start()
         # Same leader gate as the notifier above — the recurring worker jobs
         # (city warming, IndexNow) must likewise run in exactly one process.
         scheduler.start()
+        # Discord gateway bot (opt-in). Discord allows one gateway connection per
+        # bot token, so it belongs on the same single leader; it rides this event
+        # loop as a background task rather than a new process. No-op unless
+        # THERMOGRAPH_DISCORD_BOT is set with a bot token (discord_bot.enabled()).
+        if discord_bot.enabled():
+            bot_task = asyncio.create_task(discord_bot.run(), name="discord-bot")
     # Deploy/restart marker: one line correlates a metric shift with a boot.
     audit.log_activity("app.start", {"version": app.version, "role": ROLE, "base": BASE,
                                      "notifier": _should_run_notifier(), "pid": os.getpid()})
     yield
     audit.log_activity("app.stop", {"role": ROLE, "pid": os.getpid()})
+    if bot_task is not None:
+        bot_task.cancel()
+        # Wait for the cancellation to actually land before tearing anything else
+        # down -- otherwise shutdown races the bot coroutine's own teardown (closing
+        # the websocket, cancelling its heartbeat task) and the CancelledError it
+        # raises on its way out surfaces as unhandled-task noise in the logs.
+        with contextlib.suppress(asyncio.CancelledError):
+            await bot_task
     notify.stop()
     scheduler.stop()
     await _frontend_client.aclose()