From 83c2e05b967d127f37f151a2d82f7be9a79c54b4 Mon Sep 17 00:00:00 2001 From: emi Date: Thu, 23 Jul 2026 04:41:32 +0000 Subject: [PATCH] Land the Discord gateway bot (port + hardening) (#7) --- notifications/discord_bot.py | 295 ++++++++++++++++++++++++ requirements.txt | 3 + tests/notifications/test_discord_bot.py | 125 ++++++++++ web/app.py | 17 ++ 4 files changed, 440 insertions(+) create mode 100644 notifications/discord_bot.py create mode 100644 tests/notifications/test_discord_bot.py diff --git a/notifications/discord_bot.py b/notifications/discord_bot.py new file mode 100644 index 0000000..6e32780 --- /dev/null +++ b/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) diff --git a/requirements.txt b/requirements.txt index c12d96d..62fb0fc 100644 --- a/requirements.txt +++ b/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 diff --git a/tests/notifications/test_discord_bot.py b/tests/notifications/test_discord_bot.py new file mode 100644 index 0000000..9692542 --- /dev/null +++ b/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 diff --git a/web/app.py b/web/app.py index 59217ae..0d6d8bb 100644 --- a/web/app.py +++ b/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()