New app/internal/discord_rest.py: one shared httpx client with the
DiscordBot User-Agent, per-token rate-limit buckets (X-RateLimit-Bucket +
guild), one retry on a 429 of <=3s, global-429 pause per token, and a
process-wide breaker at 4000 401/403/unshared-429 responses per 10 min so
C2 can never get the shared egress IP banned. Tokens are never logged.
inspect_token / refresh_membership / leave_guild / TokenRejected move off
discord_leases onto it. /users/@me/guilds now follows after= pages (it read
one page of 200 and silently dropped the rest); a failed page is unknown,
never a truncated list.
guild_snapshot(token_doc, guild_id): cached, single-flight read of
/guilds/{g}, the bot's own member and /channels for S3's permission check.
401 marks the token invalid; 10004/10007 drop the server from guild_ids;
50001 is confirmed with one membership read and never changes guild_ids
on its own. TTLs: ok 120s, not-member/invalid 600s, unknown 30s.
config.py gets every setting the #225 slices use (perm-check kill switch,
TTLs, breaker budget, OAuth redirect URI, discord_pool_bot_guild_cap=90).
No behaviour change beyond pagination, User-Agent and the breaker.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
436 lines
19 KiB
Python
436 lines
19 KiB
Python
"""
|
|
Radio-bot token leases (server-26#204, node-26#25).
|
|
|
|
A pool "radio bot" is a Discord bot token that an edge node logs in with to
|
|
join voice channels. The rules this module enforces, all inside one Firestore
|
|
transaction per change so two concurrent requests can never both win:
|
|
|
|
* A token is leased to **at most one node** at a time (``lease.node_id``).
|
|
Discord allows one gateway session per token to be useful; two nodes logged
|
|
in on one token fight over the same voice connections.
|
|
* One leased token can serve several streams on its node, but **one per
|
|
guild** — Discord allows one voice connection per bot user per guild.
|
|
* Re-requesting a stream that already holds a lease on the same node reuses
|
|
it (``/radio play`` twice, a reconcile retry, a re-join) instead of taking a
|
|
second token.
|
|
* Release is **ordered**: stopping a stream only marks the token
|
|
``releasing``. The token returns to the pool once the node reports (MQTT
|
|
``nodes/{id}/discord/actual``) that it applied a desired-state version at
|
|
least as new as the one that dropped the token *and* has no gateway session
|
|
on it any more (``settle_releases``) — or, as a last resort, once the node
|
|
has been offline longer than the edge's own MQTT-loss fail-safe
|
|
(``release_node_leases``), at which point the edge has already logged out.
|
|
|
|
Firestore layout:
|
|
|
|
bot_tokens/{token_id}
|
|
token, name, org_id (None = platform pool, else BYO for that org),
|
|
guild_ids (list | None — None means membership never checked),
|
|
application_id, bot_user_id, valid, validated_at, check_error,
|
|
lease: None | {node_id, lease_id, since, streams: [stream_id],
|
|
releasing: bool, release_seq: int | None}
|
|
in_use / assigned_node_id / assigned_at — legacy mirrors of ``lease``,
|
|
still written so the old /tokens page and any reader of them stay right.
|
|
|
|
discord_streams/{stream_id} — see discord_streams.py. ``stream_id`` is
|
|
``{guild_id}:{channel_id}:{suffix}`` (server-26#213 follow-up: several
|
|
streams may share one voice channel, each with its own bot); docs created
|
|
before that are plain ``{guild_id}:{channel_id}`` and keep working.
|
|
|
|
Several streams in ONE guild therefore always hold DIFFERENT tokens: a token's
|
|
lease lists its streams and ``pick_token`` skips any token that already has a
|
|
stream in the guild. Both the read and the write happen inside the lease
|
|
transaction, so two concurrent starts in one guild can never share a token.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timezone
|
|
from typing import Optional
|
|
|
|
from app.internal import firestore as fstore
|
|
|
|
TOKENS = "bot_tokens"
|
|
STREAMS = "discord_streams"
|
|
|
|
# Connect (1<<20) + Speak (1<<21) + View Channel (1<<10). Nothing else — a
|
|
# radio bot never reads or sends messages, so it never asks to.
|
|
RADIO_BOT_PERMISSIONS = (1 << 10) | (1 << 20) | (1 << 21)
|
|
|
|
|
|
def invite_url(application_id: Optional[str], guild_id: Optional[str] = None) -> Optional[str]:
|
|
if not application_id:
|
|
return None
|
|
url = (
|
|
"https://discord.com/oauth2/authorize"
|
|
f"?client_id={application_id}&scope=bot&permissions={RADIO_BOT_PERMISSIONS}"
|
|
)
|
|
if guild_id:
|
|
url += f"&guild_id={guild_id}&disable_guild_select=true"
|
|
return url
|
|
|
|
|
|
# Old ids are ``guild:channel``; new ones add ``:<hex>``. Safe as a Firestore
|
|
# doc id (no "/") and as a URL path segment. Every route validates against it.
|
|
STREAM_ID_PATTERN = r"^[0-9]{1,25}:[0-9]{1,25}(?::[0-9a-f]{8,32})?$"
|
|
|
|
|
|
def stream_id_for(guild_id: str, channel_id: str) -> str:
|
|
"""The LEGACY one-stream-per-channel id. Only for reading docs created
|
|
before several streams could share a channel — never for a new stream."""
|
|
return f"{guild_id}:{channel_id}"
|
|
|
|
|
|
def new_stream_id(guild_id: str, channel_id: str) -> str:
|
|
"""A fresh id for a new stream. Unique per stream, so adding a radio
|
|
channel can never land on (and overwrite) an existing one."""
|
|
return f"{guild_id}:{channel_id}:{uuid.uuid4().hex[:12]}"
|
|
|
|
|
|
def guild_of(stream_id: str) -> str:
|
|
return stream_id.split(":", 1)[0]
|
|
|
|
|
|
def _now() -> datetime:
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
@dataclass
|
|
class LeaseResult:
|
|
ok: bool
|
|
token_id: Optional[str] = None
|
|
lease_id: Optional[str] = None
|
|
reused: bool = False
|
|
# error codes: pool_empty | no_bot_in_guild | pool_exhausted |
|
|
# preferred_busy (the chosen bot already has a stream in this guild) |
|
|
# preferred_other_node (the chosen bot is logged in on another node) |
|
|
# preferred_unavailable (unknown, invalid, or another org's bot)
|
|
error: Optional[str] = None
|
|
# bots that would work if invited to this guild: [{name, invite_url}]
|
|
invite: list = field(default_factory=list)
|
|
# preferred_* errors: {token_name, busy_stream_id, node_id}
|
|
detail: dict = field(default_factory=dict)
|
|
|
|
|
|
def _lease_streams(tok: dict) -> list:
|
|
return list(((tok.get("lease") or {}).get("streams")) or [])
|
|
|
|
|
|
def _legacy_mirror(lease: Optional[dict]) -> dict:
|
|
if not lease:
|
|
return {"lease": None, "in_use": False, "assigned_node_id": None, "assigned_at": None}
|
|
return {
|
|
"lease": lease,
|
|
"in_use": True,
|
|
"assigned_node_id": lease["node_id"],
|
|
"assigned_at": lease.get("since"),
|
|
}
|
|
|
|
|
|
def _candidate_order(tok_id: str, tok: dict, node_id: str, org_id: Optional[str],
|
|
preferred: Optional[str]) -> tuple:
|
|
lease = tok.get("lease") or {}
|
|
return (
|
|
0 if tok_id == preferred else 1,
|
|
0 if org_id and tok.get("org_id") == org_id else 1, # BYO before pool
|
|
0 if lease.get("node_id") == node_id else 1, # reuse a session
|
|
0 if tok.get("guild_ids") is not None else 1, # known member first
|
|
str(tok.get("name") or ""),
|
|
tok_id,
|
|
)
|
|
|
|
|
|
def pick_token(
|
|
tokens: dict[str, dict],
|
|
*,
|
|
node_id: str,
|
|
guild_id: str,
|
|
org_id: Optional[str],
|
|
preferred_token_id: Optional[str] = None,
|
|
exclude_stream_id: Optional[str] = None,
|
|
) -> LeaseResult:
|
|
"""Pure choice over a snapshot of bot_tokens. No I/O — unit tested directly.
|
|
|
|
A ``preferred_token_id`` is a demand, not a hint: if that bot cannot take
|
|
the stream the result says why instead of quietly using another bot."""
|
|
if not tokens:
|
|
return LeaseResult(ok=False, error="pool_empty")
|
|
if preferred_token_id:
|
|
return _pick_preferred(tokens, preferred_token_id, node_id=node_id, guild_id=guild_id,
|
|
org_id=org_id, exclude_stream_id=exclude_stream_id)
|
|
usable: list[tuple] = []
|
|
not_member: list[dict] = []
|
|
for tok_id, tok in tokens.items():
|
|
if tok.get("valid") is False:
|
|
continue
|
|
owner = tok.get("org_id")
|
|
if owner is not None and owner != org_id:
|
|
continue # another org's branded bot — never leased across orgs
|
|
lease = tok.get("lease") or {}
|
|
if lease and lease.get("node_id") and lease.get("node_id") != node_id:
|
|
continue # live (or still releasing) on another node
|
|
busy_guilds = {guild_of(s) for s in _lease_streams(tok) if s != exclude_stream_id}
|
|
if guild_id in busy_guilds:
|
|
continue # this bot already has a voice connection in this guild
|
|
members = tok.get("guild_ids")
|
|
if members is not None and guild_id not in members:
|
|
not_member.append({"name": tok.get("name") or tok_id,
|
|
"invite_url": invite_url(tok.get("application_id"), guild_id)})
|
|
continue
|
|
usable.append((_candidate_order(tok_id, tok, node_id, org_id, preferred_token_id), tok_id))
|
|
if usable:
|
|
usable.sort()
|
|
return LeaseResult(ok=True, token_id=usable[0][1])
|
|
if not_member:
|
|
return LeaseResult(ok=False, error="no_bot_in_guild", invite=not_member)
|
|
return LeaseResult(ok=False, error="pool_exhausted")
|
|
|
|
|
|
def _pick_preferred(tokens: dict[str, dict], tok_id: str, *, node_id: str, guild_id: str,
|
|
org_id: Optional[str], exclude_stream_id: Optional[str]) -> LeaseResult:
|
|
tok = tokens.get(tok_id)
|
|
owner = (tok or {}).get("org_id")
|
|
if not tok or tok.get("valid") is False or (owner is not None and owner != org_id):
|
|
return LeaseResult(ok=False, error="preferred_unavailable",
|
|
detail={"token_name": (tok or {}).get("name") or tok_id})
|
|
name = tok.get("name") or tok_id
|
|
busy = [s for s in _lease_streams(tok) if s != exclude_stream_id and guild_of(s) == guild_id]
|
|
if busy:
|
|
# One voice connection per bot per guild (Discord): never move the
|
|
# bot off another stream to serve this one.
|
|
return LeaseResult(ok=False, error="preferred_busy",
|
|
detail={"token_name": name, "busy_stream_id": busy[0]})
|
|
lease = tok.get("lease") or {}
|
|
if lease.get("node_id") and lease["node_id"] != node_id:
|
|
return LeaseResult(ok=False, error="preferred_other_node",
|
|
detail={"token_name": name, "node_id": lease["node_id"]})
|
|
members = tok.get("guild_ids")
|
|
if members is not None and guild_id not in members:
|
|
return LeaseResult(ok=False, error="no_bot_in_guild",
|
|
invite=[{"name": name, "invite_url": invite_url(tok.get("application_id"), guild_id)}])
|
|
return LeaseResult(ok=True, token_id=tok_id)
|
|
|
|
|
|
async def lease_for_stream(
|
|
stream_id: str,
|
|
*,
|
|
node_id: str,
|
|
org_id: Optional[str],
|
|
preferred_token_id: Optional[str] = None,
|
|
stream_fields: Optional[dict] = None,
|
|
) -> LeaseResult:
|
|
"""Bind ``stream_id`` to a token leased to ``node_id``, atomically.
|
|
|
|
Also writes ``stream_fields`` (desired=on, node_id, token_id, lease_id...)
|
|
onto the stream doc in the same transaction, so a stream doc can never name
|
|
a token whose lease does not list it.
|
|
"""
|
|
guild_id = guild_of(stream_id)
|
|
db = fstore.db
|
|
|
|
def _txn(txn):
|
|
stream_ref = db.collection(STREAMS).document(stream_id)
|
|
stream_snap = stream_ref.get(transaction=txn)
|
|
stream = stream_snap.to_dict() if stream_snap.exists else {}
|
|
tokens = {s.id: s.to_dict() for s in db.collection(TOKENS).stream(transaction=txn)}
|
|
|
|
prev_tok_id = stream.get("token_id")
|
|
prev_tok = tokens.get(prev_tok_id) if prev_tok_id else None
|
|
if (
|
|
prev_tok
|
|
and stream.get("node_id") == node_id
|
|
and (prev_tok.get("lease") or {}).get("node_id") == node_id
|
|
and stream_id in _lease_streams(prev_tok)
|
|
and prev_tok.get("valid") is not False
|
|
# An explicit edit to a different bot moves the stream to it.
|
|
and (not preferred_token_id or preferred_token_id == prev_tok_id)
|
|
):
|
|
lease = prev_tok["lease"]
|
|
if lease.get("releasing"):
|
|
lease = {**lease, "releasing": False, "release_seq": None}
|
|
txn.update(db.collection(TOKENS).document(prev_tok_id), _legacy_mirror(lease))
|
|
txn.set(stream_ref, {**(stream_fields or {}), "node_id": node_id,
|
|
"token_id": prev_tok_id, "lease_id": lease["lease_id"]}, merge=True)
|
|
return LeaseResult(ok=True, token_id=prev_tok_id, lease_id=lease["lease_id"], reused=True)
|
|
|
|
result = pick_token(tokens, node_id=node_id, guild_id=guild_id, org_id=org_id,
|
|
preferred_token_id=preferred_token_id, exclude_stream_id=stream_id)
|
|
if not result.ok:
|
|
return result
|
|
|
|
# Detach from the previous token (moved node, or token went invalid).
|
|
if prev_tok and stream_id in _lease_streams(prev_tok) and prev_tok_id != result.token_id:
|
|
txn.update(db.collection(TOKENS).document(prev_tok_id),
|
|
_legacy_mirror(_without_stream(prev_tok["lease"], stream_id)))
|
|
|
|
tok = tokens[result.token_id]
|
|
lease = tok.get("lease") or {}
|
|
if lease.get("node_id") == node_id:
|
|
streams = [s for s in _lease_streams(tok) if s != stream_id] + [stream_id]
|
|
lease = {**lease, "streams": streams, "releasing": False, "release_seq": None}
|
|
else:
|
|
lease = {"node_id": node_id, "lease_id": uuid.uuid4().hex, "since": _now().isoformat(),
|
|
"streams": [stream_id], "releasing": False, "release_seq": None}
|
|
txn.update(db.collection(TOKENS).document(result.token_id), _legacy_mirror(lease))
|
|
txn.set(stream_ref, {**(stream_fields or {}), "node_id": node_id,
|
|
"token_id": result.token_id, "lease_id": lease["lease_id"]}, merge=True)
|
|
result.lease_id = lease["lease_id"]
|
|
return result
|
|
|
|
return await fstore.run_transaction(_txn)
|
|
|
|
|
|
def _without_stream(lease: dict, stream_id: str) -> dict:
|
|
streams = [s for s in (lease.get("streams") or []) if s != stream_id]
|
|
# The token stays leased to the node — it may still be logged in there —
|
|
# until settle_releases sees the node ack. ``releasing`` only once empty.
|
|
return {**lease, "streams": streams, "releasing": not streams,
|
|
"release_seq": lease.get("release_seq") if streams else None}
|
|
|
|
|
|
async def detach_stream(stream_id: str, stream_updates: Optional[dict] = None) -> Optional[str]:
|
|
"""Drop ``stream_id`` from its token's lease (token -> ``releasing`` if that
|
|
was its last stream) and apply ``stream_updates`` to the stream doc, in one
|
|
transaction. Returns the node_id the token is (still) leased to, if any."""
|
|
db = fstore.db
|
|
|
|
def _txn(txn):
|
|
stream_ref = db.collection(STREAMS).document(stream_id)
|
|
snap = stream_ref.get(transaction=txn)
|
|
if not snap.exists:
|
|
return None
|
|
stream = snap.to_dict()
|
|
tok_id = stream.get("token_id")
|
|
tok_node = None
|
|
if tok_id:
|
|
tok_ref = db.collection(TOKENS).document(tok_id)
|
|
tsnap = tok_ref.get(transaction=txn)
|
|
tok = tsnap.to_dict() if tsnap.exists else None
|
|
if tok and tok.get("lease") and stream_id in _lease_streams(tok):
|
|
txn.update(tok_ref, _legacy_mirror(_without_stream(tok["lease"], stream_id)))
|
|
tok_node = tok["lease"].get("node_id")
|
|
txn.set(stream_ref, {**(stream_updates or {}), "token_id": None, "lease_id": None}, merge=True)
|
|
return tok_node or stream.get("node_id")
|
|
|
|
return await fstore.run_transaction(_txn)
|
|
|
|
|
|
async def delete_stopped_stream(stream_id: str) -> bool:
|
|
"""Delete a stream doc, but only while it is stopped and holds no token.
|
|
|
|
True when the doc is gone (deleted now, or already absent). False when it
|
|
is desired-on again or names a token — a concurrent start re-leased it,
|
|
and deleting it would leave that token's lease listing a stream nobody can
|
|
stop. ``lease_for_stream`` writes the stream doc in its own transaction,
|
|
so reading it here is enough to serialise against it."""
|
|
db = fstore.db
|
|
|
|
def _txn(txn):
|
|
ref = db.collection(STREAMS).document(stream_id)
|
|
snap = ref.get(transaction=txn)
|
|
if not snap.exists:
|
|
return True
|
|
stream = snap.to_dict() or {}
|
|
if stream.get("desired") == "on" or stream.get("token_id"):
|
|
return False
|
|
txn.delete(ref)
|
|
return True
|
|
|
|
return await fstore.run_transaction(_txn)
|
|
|
|
|
|
async def stamp_release_seq(node_id: str, seq: int) -> None:
|
|
"""Record the desired-state version that dropped each releasing token, so a
|
|
stale actual report (from before the node saw the drop) can never free it."""
|
|
db = fstore.db
|
|
|
|
def _txn(txn):
|
|
tokens = [(s.id, s.to_dict()) for s in db.collection(TOKENS).stream(transaction=txn)]
|
|
for tok_id, tok in tokens:
|
|
lease = tok.get("lease") or {}
|
|
if lease.get("node_id") == node_id and lease.get("releasing") and lease.get("release_seq") is None:
|
|
txn.update(db.collection(TOKENS).document(tok_id),
|
|
_legacy_mirror({**lease, "release_seq": seq}))
|
|
|
|
await fstore.run_transaction(_txn)
|
|
|
|
|
|
async def settle_releases(node_id: str, sessions: set, applied_seq: Optional[int]) -> list[str]:
|
|
"""Free tokens the node has confirmed it no longer runs. Returns freed ids."""
|
|
db = fstore.db
|
|
|
|
def _txn(txn):
|
|
freed = []
|
|
tokens = [(s.id, s.to_dict()) for s in db.collection(TOKENS).stream(transaction=txn)]
|
|
for tok_id, tok in tokens:
|
|
lease = tok.get("lease") or {}
|
|
if lease.get("node_id") != node_id or not lease.get("releasing") or lease.get("streams"):
|
|
continue
|
|
rseq = lease.get("release_seq")
|
|
if rseq is None or applied_seq is None or applied_seq < rseq:
|
|
continue # the node may not have seen the drop yet
|
|
if tok_id in sessions:
|
|
continue # still logged in on the node
|
|
txn.update(db.collection(TOKENS).document(tok_id), _legacy_mirror(None))
|
|
freed.append(tok_id)
|
|
return freed
|
|
|
|
return await fstore.run_transaction(_txn)
|
|
|
|
|
|
async def release_node_leases(node_id: str) -> list[str]:
|
|
"""Last-resort release for a node that has been offline longer than the
|
|
edge's MQTT-loss fail-safe (the edge logs every radio bot out by itself
|
|
after DISCORD_MQTT_LOSS_STOP_SECONDS without the broker). Returns the
|
|
stream ids that lost their token so the caller can re-place them."""
|
|
db = fstore.db
|
|
|
|
def _txn(txn):
|
|
orphaned = []
|
|
tokens = [(s.id, s.to_dict()) for s in db.collection(TOKENS).stream(transaction=txn)]
|
|
for tok_id, tok in tokens:
|
|
lease = tok.get("lease") or {}
|
|
if lease.get("node_id") != node_id:
|
|
continue
|
|
orphaned.extend(lease.get("streams") or [])
|
|
txn.update(db.collection(TOKENS).document(tok_id), _legacy_mirror(None))
|
|
return orphaned
|
|
|
|
return await fstore.run_transaction(_txn)
|
|
|
|
|
|
async def force_release_token(token_id: str) -> Optional[dict]:
|
|
"""Admin override (/tokens/flush, delete). Frees the token and detaches
|
|
every stream that named it, in one transaction, so no stream doc is left
|
|
pointing at a token whose lease no longer lists it. Returns the old lease
|
|
(``None`` if the token was not leased or does not exist); the caller
|
|
republishes desired state for ``lease["node_id"]`` so the edge logs out."""
|
|
db = fstore.db
|
|
|
|
def _txn(txn):
|
|
ref = db.collection(TOKENS).document(token_id)
|
|
snap = ref.get(transaction=txn)
|
|
if not snap.exists:
|
|
return None
|
|
lease = snap.to_dict().get("lease")
|
|
stream_refs = []
|
|
for sid in (lease or {}).get("streams") or []:
|
|
sref = db.collection(STREAMS).document(sid)
|
|
ssnap = sref.get(transaction=txn)
|
|
if ssnap.exists and (ssnap.to_dict() or {}).get("token_id") == token_id:
|
|
stream_refs.append(sref)
|
|
txn.update(ref, _legacy_mirror(None))
|
|
for sref in stream_refs:
|
|
txn.set(sref, {"token_id": None, "lease_id": None}, merge=True)
|
|
return lease
|
|
|
|
return await fstore.run_transaction(_txn)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Token validation / guild membership / leaving a server moved to
|
|
# app/internal/discord_rest.py (server-26#225 S2): inspect_token,
|
|
# refresh_membership, leave_guild, TokenRejected.
|
|
# ---------------------------------------------------------------------------
|