""" 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 ``:``. 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 (server-26#204 item 5) # --------------------------------------------------------------------------- DISCORD_API = "https://discord.com/api/v10" class TokenRejected(Exception): """Discord answered 401: the string is not a working bot token.""" async def inspect_token(token: str) -> dict: """Ask Discord who this token is and which servers it is in. Returns ``{bot_user_id, bot_name, application_id, guild_ids, valid, validated_at, check_error}``. Raises TokenRejected on a 401. Any other failure (Discord down, rate limit) is reported in ``check_error`` with ``valid``/``guild_ids`` = None — unknown, not bad: a token is never rejected because Discord had a bad minute. """ import httpx headers = {"Authorization": f"Bot {token}"} out: dict = {"valid": None, "guild_ids": None, "check_error": None, "validated_at": _now().isoformat()} try: async with httpx.AsyncClient(timeout=10, base_url=DISCORD_API, headers=headers) as client: me = await client.get("/users/@me") if me.status_code == 401: raise TokenRejected("Discord rejected this bot token (401).") me.raise_for_status() user = me.json() out["bot_user_id"] = str(user.get("id")) out["bot_name"] = user.get("username") app = await client.get("/oauth2/applications/@me") if app.status_code == 200: out["application_id"] = str(app.json().get("id")) guilds = await client.get("/users/@me/guilds", params={"limit": 200}) guilds.raise_for_status() out["guild_ids"] = [str(g["id"]) for g in guilds.json()] out["valid"] = True except TokenRejected: raise except Exception as e: # network / 5xx / 429 — unknown, not invalid out["check_error"] = str(e)[:200] return out async def refresh_membership(token_ids: list[str]) -> int: """Re-read guild membership for these tokens. Returns how many changed.""" changed = 0 for tok_id in token_ids: tok = await fstore.doc_get(TOKENS, tok_id) if not tok or not tok.get("token"): continue try: info = await inspect_token(tok["token"]) except TokenRejected as e: await fstore.doc_update(TOKENS, tok_id, {"valid": False, "check_error": str(e), "validated_at": _now().isoformat()}) changed += 1 continue updates = {k: v for k, v in info.items() if v is not None or k == "check_error"} if updates.get("guild_ids") != tok.get("guild_ids") or updates.get("valid") != tok.get("valid"): changed += 1 await fstore.doc_update(TOKENS, tok_id, updates) return changed async def leave_guild(guild_id: str, org_id: Optional[str]) -> list[dict]: """Make the radio bots DRB holds tokens for leave one Discord server (Discord REST ``DELETE /users/@me/guilds/{id}``, the bot's own token). Only platform-pool bots and ``org_id``'s own BYO bots — never another org's branded bot. Bots whose cached membership says they are not in the server are skipped; unknown membership is tried (404 = not a member). The DRB *server* bot (the discord-bot service) is not in this pool, so it stays in the server until someone kicks it. Returns one row per bot tried: ``{token_id, name, left, error}``. Never raises for a Discord failure.""" import httpx guild_id = str(guild_id) out: list[dict] = [] for tok in await fstore.collection_list(TOKENS): tok_id = tok.get("token_id") if not tok_id or not tok.get("token"): continue if tok.get("org_id") is not None and tok.get("org_id") != org_id: continue members = tok.get("guild_ids") if members is not None and guild_id not in members: continue row = {"token_id": tok_id, "name": tok.get("name") or tok_id, "left": False, "error": None} try: async with httpx.AsyncClient(timeout=10, base_url=DISCORD_API, headers={"Authorization": f"Bot {tok['token']}"}) as client: r = await client.delete(f"/users/@me/guilds/{guild_id}") if r.status_code in (200, 204, 404): # 404: Unknown Guild — not a member row["left"] = True else: row["error"] = f"Discord answered {r.status_code}" except Exception as e: # network — report, do not fail the removal row["error"] = str(e)[:200] if row["left"] and members is not None: await fstore.doc_update(TOKENS, tok_id, {"guild_ids": [g for g in members if g != guild_id]}) out.append(row) return out