Logan CusanoandClaude Opus 5.5
0a724a064f
Paused by owner instruction; not wired into routers/mqtt yet, untested. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
334 lines
14 KiB
Python
334 lines
14 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/{guild_id}:{channel_id} — see discord_streams.py.
|
|
"""
|
|
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
|
|
|
|
|
|
def stream_id_for(guild_id: str, channel_id: str) -> str:
|
|
return f"{guild_id}:{channel_id}"
|
|
|
|
|
|
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
|
|
error: Optional[str] = None
|
|
# bots that would work if invited to this guild: [{name, invite_url}]
|
|
invite: list = field(default_factory=list)
|
|
|
|
|
|
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."""
|
|
if not tokens:
|
|
return LeaseResult(ok=False, error="pool_empty")
|
|
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")
|
|
|
|
|
|
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
|
|
):
|
|
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 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, Force release). Returns the old lease."""
|
|
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")
|
|
txn.update(ref, _legacy_mirror(None))
|
|
return lease
|
|
|
|
return await fstore.run_transaction(_txn)
|