295 lines
13 KiB
Python
295 lines
13 KiB
Python
"""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)
|