WIP: discord system — transactional token leases + desired-state streams (c2-core, unwired)
Paused by owner instruction; not wired into routers/mqtt yet, untested. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Logan Cusano
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
d57738a8f6
commit
0a724a064f
@@ -0,0 +1,333 @@
|
||||
"""
|
||||
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)
|
||||
@@ -0,0 +1,376 @@
|
||||
"""
|
||||
Discord radio channels as desired state (server-26#203 MVP).
|
||||
|
||||
One ``discord_streams/{guild_id}:{channel_id}`` doc per Discord voice channel
|
||||
that should carry radio. Fields:
|
||||
|
||||
org_id, guild_id, channel_id, channel_name, guild_name,
|
||||
system_id, system_name,
|
||||
mode: "persistent" | "adhoc"
|
||||
persistent — stays connected and reconnects by itself (node reboot,
|
||||
Discord voice drop, C2 restart) until someone stops it.
|
||||
adhoc — the Discord bot sets desired=off after ``idle_minutes``
|
||||
with no humans in the voice channel.
|
||||
desired: "on" | "off"
|
||||
node_id, token_id, lease_id — placement (written with the lease, atomically)
|
||||
actual: {state, error, since, at} — reported by the node, never guessed here
|
||||
state: pending | connecting | connected | error | stopped | node_offline
|
||||
tg_filter: None — phase-2 seam (per-channel talkgroup gate); always None now
|
||||
now_playing_message_id, created_by_discord_id, created_at, updated_at
|
||||
|
||||
Flow: ``start_stream`` picks a node + leases a token -> ``publish_desired``
|
||||
retains ``nodes/{node}/discord/desired`` (full list for that node, with a
|
||||
monotonically increasing ``seq``) -> the edge reconciles and publishes
|
||||
``nodes/{node}/discord/actual`` -> ``handle_actual`` updates ``actual`` and
|
||||
settles token releases. ``reconcile_loop`` re-places persistent streams whose
|
||||
node went away.
|
||||
|
||||
The desired payload carries token **ids**, never token strings: the edge
|
||||
fetches the string over HTTPS (GET /discord/node/tokens/{id}) with its own
|
||||
node API key, and C2 hands it out only to the node holding the lease. A
|
||||
retained MQTT message therefore never stores a bot credential on the broker.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from datetime import datetime, timezone, timedelta
|
||||
from typing import Optional
|
||||
|
||||
from app.config import settings
|
||||
from app.internal import firestore as fstore
|
||||
from app.internal import discord_leases as leases
|
||||
from app.internal.logger import logger
|
||||
|
||||
STREAMS = leases.STREAMS
|
||||
NODE_STATE = "discord_nodes" # discord_nodes/{node_id}: {seq, applied_seq, actual_at}
|
||||
ONLINE = ("online", "recording")
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _effective_system(node: dict) -> Optional[str]:
|
||||
if node.get("is_overridden") and node.get("override_system_id"):
|
||||
return node["override_system_id"]
|
||||
return node.get("assigned_system_id")
|
||||
|
||||
|
||||
async def _node_stream_counts(org_id: Optional[str]) -> dict[str, int]:
|
||||
streams = await fstore.collection_list(STREAMS, desired="on")
|
||||
counts: dict[str, int] = {}
|
||||
for s in streams:
|
||||
if s.get("node_id"):
|
||||
counts[s["node_id"]] = counts.get(s["node_id"], 0) + 1
|
||||
return counts
|
||||
|
||||
|
||||
async def candidate_nodes(org_id: str, system_id: str, exclude: Optional[str] = None) -> list[dict]:
|
||||
"""Online, approved nodes of this org currently listening to ``system_id``,
|
||||
best first: nodes already carrying streams (one gateway session can serve
|
||||
several guilds) but under the per-node cap, then the rest."""
|
||||
nodes = await fstore.collection_list("nodes", org_id=org_id)
|
||||
counts = await _node_stream_counts(org_id)
|
||||
cap = settings.discord_node_max_streams
|
||||
out = []
|
||||
for n in nodes:
|
||||
if n.get("node_id") == exclude:
|
||||
continue
|
||||
if n.get("status") not in ONLINE or n.get("approval_status") not in (None, "approved"):
|
||||
continue
|
||||
if _effective_system(n) != system_id:
|
||||
continue
|
||||
if counts.get(n["node_id"], 0) >= cap:
|
||||
continue
|
||||
out.append(n)
|
||||
out.sort(key=lambda n: (-counts.get(n["node_id"], 0), n.get("name") or n["node_id"]))
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Desired state -> node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def _bump_seq(node_id: str) -> int:
|
||||
db = fstore.db
|
||||
|
||||
def _txn(txn):
|
||||
ref = db.collection(NODE_STATE).document(node_id)
|
||||
snap = ref.get(transaction=txn)
|
||||
seq = int((snap.to_dict() or {}).get("seq") or 0) + 1 if snap.exists else 1
|
||||
txn.set(ref, {"node_id": node_id, "seq": seq}, merge=True)
|
||||
return seq
|
||||
|
||||
return await fstore.run_transaction(_txn)
|
||||
|
||||
|
||||
async def desired_for_node(node_id: str) -> list[dict]:
|
||||
streams = await fstore.collection_list(STREAMS, node_id=node_id)
|
||||
out = []
|
||||
for s in streams:
|
||||
if s.get("desired") != "on" or not s.get("token_id"):
|
||||
continue
|
||||
out.append({
|
||||
"stream_id": s["stream_id"],
|
||||
"guild_id": s["guild_id"],
|
||||
"channel_id": s["channel_id"],
|
||||
"token_id": s["token_id"],
|
||||
"lease_id": s.get("lease_id"),
|
||||
"system_name": s.get("system_name"),
|
||||
"mode": s.get("mode", "adhoc"),
|
||||
# Phase-2 seam: the edge ignores this today (no TG gate yet).
|
||||
"tg_filter": s.get("tg_filter"),
|
||||
})
|
||||
out.sort(key=lambda d: d["stream_id"])
|
||||
return out
|
||||
|
||||
|
||||
async def publish_desired(node_id: str) -> int:
|
||||
"""Retain the node's full stream list; stamp releasing tokens with the seq."""
|
||||
from app.internal.mqtt_handler import mqtt_handler
|
||||
seq = await _bump_seq(node_id)
|
||||
payload = {"seq": seq, "streams": await desired_for_node(node_id), "issued_at": _now_iso()}
|
||||
mqtt_handler.publish_discord_desired(node_id, payload)
|
||||
await leases.stamp_release_seq(node_id, seq)
|
||||
return seq
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Start / stop
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class StreamError(Exception):
|
||||
def __init__(self, code: str, message: str, invite: Optional[list] = None, status: int = 409):
|
||||
super().__init__(message)
|
||||
self.code = code
|
||||
self.message = message
|
||||
self.invite = invite or []
|
||||
self.status = status
|
||||
|
||||
|
||||
ERROR_TEXT = {
|
||||
"no_node": "No online node in your organization is listening to that system right now.",
|
||||
"pool_empty": "No radio bots have been added to DRB yet. A platform admin adds them under Radio bots.",
|
||||
"no_bot_in_guild": "No free radio bot is a member of this server. Invite one with the link below, then try again.",
|
||||
"pool_exhausted": "Every radio bot that can join this server is already busy. Stop another radio channel or add a bot.",
|
||||
"limit": "This server already has its maximum number of radio channels. Stop one with /radio stop first.",
|
||||
}
|
||||
|
||||
|
||||
async def start_stream(
|
||||
*,
|
||||
org_id: str,
|
||||
guild_id: str,
|
||||
channel_id: str,
|
||||
system_id: str,
|
||||
mode: str = "adhoc",
|
||||
channel_name: Optional[str] = None,
|
||||
guild_name: Optional[str] = None,
|
||||
created_by: Optional[str] = None,
|
||||
node_id: Optional[str] = None,
|
||||
preferred_token_id: Optional[str] = None,
|
||||
max_streams: Optional[int] = None,
|
||||
) -> dict:
|
||||
"""Make ``guild_id:channel_id`` carry ``system_id``. Idempotent: re-running
|
||||
it for a channel that is already on reuses the node and the lease."""
|
||||
if mode not in ("persistent", "adhoc"):
|
||||
raise StreamError("bad_mode", "mode must be persistent or adhoc", status=400)
|
||||
stream_id = leases.stream_id_for(guild_id, channel_id)
|
||||
existing = await fstore.doc_get(STREAMS, stream_id) or {}
|
||||
|
||||
if max_streams is not None and existing.get("desired") != "on":
|
||||
live = [s for s in await fstore.collection_list(STREAMS, guild_id=guild_id) if s.get("desired") == "on"]
|
||||
if len(live) >= max_streams:
|
||||
raise StreamError("limit", ERROR_TEXT["limit"])
|
||||
|
||||
system = await fstore.doc_get_cached("systems", system_id) or {}
|
||||
|
||||
if node_id:
|
||||
candidates = [await fstore.doc_get("nodes", node_id) or {"node_id": node_id}]
|
||||
else:
|
||||
candidates = await candidate_nodes(org_id, system_id)
|
||||
# Keep the current node when it is still a valid choice (reuse the lease).
|
||||
cur = existing.get("node_id")
|
||||
if cur and existing.get("system_id") == system_id:
|
||||
candidates.sort(key=lambda n: 0 if n["node_id"] == cur else 1)
|
||||
if not candidates:
|
||||
raise StreamError("no_node", ERROR_TEXT["no_node"])
|
||||
|
||||
now = _now_iso()
|
||||
fields = {
|
||||
"stream_id": stream_id, "org_id": org_id, "guild_id": guild_id, "channel_id": channel_id,
|
||||
"system_id": system_id, "system_name": system.get("name") or system_id,
|
||||
"mode": mode, "desired": "on", "tg_filter": None,
|
||||
"actual": {"state": "pending", "error": None, "since": now, "at": now},
|
||||
"updated_at": now,
|
||||
}
|
||||
if channel_name:
|
||||
fields["channel_name"] = channel_name
|
||||
if guild_name:
|
||||
fields["guild_name"] = guild_name
|
||||
if not existing:
|
||||
fields["created_at"] = now
|
||||
fields["created_by_discord_id"] = created_by
|
||||
fields["now_playing_message_id"] = None
|
||||
|
||||
last: Optional[leases.LeaseResult] = None
|
||||
prev_node = existing.get("node_id")
|
||||
for node in candidates:
|
||||
res = await leases.lease_for_stream(stream_id, node_id=node["node_id"], org_id=org_id,
|
||||
preferred_token_id=preferred_token_id, stream_fields=fields)
|
||||
if res.ok:
|
||||
await publish_desired(node["node_id"])
|
||||
if prev_node and prev_node != node["node_id"]:
|
||||
await publish_desired(prev_node)
|
||||
logger.info(f"Discord stream {stream_id} -> node {node['node_id']} token {res.token_id}"
|
||||
f" ({'reused' if res.reused else 'new'} lease)")
|
||||
return {**existing, **fields, "node_id": node["node_id"], "node_name": node.get("name"),
|
||||
"token_id": res.token_id, "lease_id": res.lease_id, "reused": res.reused}
|
||||
last = res
|
||||
code = last.error if last else "no_node"
|
||||
raise StreamError(code, ERROR_TEXT.get(code, code), invite=last.invite if last else [])
|
||||
|
||||
|
||||
async def stop_stream(stream_id: str, reason: str = "stopped") -> bool:
|
||||
existing = await fstore.doc_get(STREAMS, stream_id)
|
||||
if not existing:
|
||||
return False
|
||||
now = _now_iso()
|
||||
node_id = await leases.detach_stream(stream_id, {
|
||||
"desired": "off", "updated_at": now, "stop_reason": reason,
|
||||
"actual": {**(existing.get("actual") or {}), "state": "stopping", "at": now},
|
||||
})
|
||||
if node_id:
|
||||
await publish_desired(node_id)
|
||||
logger.info(f"Discord stream {stream_id} stopped ({reason})")
|
||||
return True
|
||||
|
||||
|
||||
async def stop_guild(guild_id: str, reason: str) -> int:
|
||||
n = 0
|
||||
for s in await fstore.collection_list(STREAMS, guild_id=guild_id):
|
||||
if s.get("desired") == "on":
|
||||
await stop_stream(s["stream_id"], reason)
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
async def stop_node(node_id: str, reason: str) -> int:
|
||||
n = 0
|
||||
for s in await fstore.collection_list(STREAMS, node_id=node_id):
|
||||
if s.get("desired") == "on":
|
||||
await stop_stream(s["stream_id"], reason)
|
||||
n += 1
|
||||
return n
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Actual state <- node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_VALID_STATES = {"connecting", "connected", "error", "stopped", "paused_locally"}
|
||||
|
||||
|
||||
async def handle_actual(node_id: str, payload: dict) -> None:
|
||||
"""``nodes/{node_id}/discord/actual``. A node may only report on streams
|
||||
currently placed on it — anything else is ignored, so a compromised node
|
||||
cannot rewrite another node's (or org's) stream state."""
|
||||
reported = payload.get("streams") or {}
|
||||
now = _now_iso()
|
||||
if isinstance(reported, dict):
|
||||
for stream_id, st in list(reported.items())[:50]:
|
||||
if not isinstance(st, dict):
|
||||
continue
|
||||
doc = await fstore.doc_get(STREAMS, str(stream_id))
|
||||
if not doc or doc.get("node_id") != node_id or doc.get("desired") != "on":
|
||||
continue
|
||||
state = st.get("state") if st.get("state") in _VALID_STATES else "error"
|
||||
error = st.get("error")
|
||||
error = str(error)[:300] if error else None
|
||||
prev = doc.get("actual") or {}
|
||||
since = prev.get("since") if prev.get("state") == state else now
|
||||
await fstore.doc_update(STREAMS, stream_id, {
|
||||
"actual": {"state": state, "error": error, "since": since, "at": now},
|
||||
})
|
||||
|
||||
if "sessions" in payload:
|
||||
sessions = {str(t) for t in (payload.get("sessions") or [])}
|
||||
applied = payload.get("applied_seq")
|
||||
applied = int(applied) if isinstance(applied, (int, float)) else None
|
||||
freed = await leases.settle_releases(node_id, sessions, applied)
|
||||
for tok in freed:
|
||||
logger.info(f"Radio bot token {tok} released — node {node_id} confirmed it logged out")
|
||||
await fstore.doc_set(NODE_STATE, node_id, {"applied_seq": applied, "actual_at": now})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Self-healing: re-place persistent streams whose node is gone
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def release_offline_node(node_id: str) -> None:
|
||||
"""Node offline past the lease grace: free its tokens (the edge already
|
||||
logged them out — see release_node_leases) and re-place its streams."""
|
||||
orphaned = await leases.release_node_leases(node_id)
|
||||
if not orphaned:
|
||||
return
|
||||
now = _now_iso()
|
||||
for sid in orphaned:
|
||||
await fstore.doc_set(STREAMS, sid, {
|
||||
"token_id": None, "lease_id": None,
|
||||
"actual": {"state": "node_offline", "error": "The node carrying this channel went offline.",
|
||||
"since": now, "at": now},
|
||||
})
|
||||
await publish_desired(node_id) # retained: the node sees an accurate (empty) list when it returns
|
||||
logger.warning(f"Node {node_id} offline past lease grace — released {len(orphaned)} radio stream(s)")
|
||||
|
||||
|
||||
async def reconcile_once() -> None:
|
||||
nodes = {n["node_id"]: n for n in await fstore.collection_list("nodes") if n.get("node_id")}
|
||||
grace = timedelta(seconds=settings.discord_lease_offline_release_seconds)
|
||||
now = datetime.now(timezone.utc)
|
||||
for s in await fstore.collection_list(STREAMS, desired="on"):
|
||||
node = nodes.get(s.get("node_id") or "")
|
||||
online = bool(node) and node.get("status") in ONLINE and _effective_system(node) == s.get("system_id")
|
||||
if online and s.get("token_id"):
|
||||
continue
|
||||
if node and s.get("token_id"):
|
||||
# Node offline (or retuned): hold the lease until the grace passes.
|
||||
last_seen = node.get("last_seen")
|
||||
try:
|
||||
ls = datetime.fromisoformat(last_seen) if isinstance(last_seen, str) else last_seen
|
||||
if ls and ls.tzinfo is None:
|
||||
ls = ls.replace(tzinfo=timezone.utc)
|
||||
except Exception:
|
||||
ls = None
|
||||
if node.get("status") in ONLINE or (ls and now - ls < grace):
|
||||
if node.get("status") not in ONLINE:
|
||||
await _mark(s, "node_offline", "The node carrying this channel is offline; waiting for it.")
|
||||
continue
|
||||
await release_offline_node(node["node_id"])
|
||||
if s.get("mode") != "persistent":
|
||||
await stop_stream(s["stream_id"], "node_lost")
|
||||
continue
|
||||
try:
|
||||
await start_stream(org_id=s["org_id"], guild_id=s["guild_id"], channel_id=s["channel_id"],
|
||||
system_id=s["system_id"], mode="persistent")
|
||||
except StreamError as e:
|
||||
await _mark(s, "error", e.message)
|
||||
|
||||
|
||||
async def _mark(s: dict, state: str, error: str) -> None:
|
||||
prev = s.get("actual") or {}
|
||||
if prev.get("state") == state and prev.get("error") == error:
|
||||
return
|
||||
now = _now_iso()
|
||||
await fstore.doc_set(STREAMS, s["stream_id"], {"actual": {"state": state, "error": error,
|
||||
"since": now, "at": now}})
|
||||
|
||||
|
||||
async def reconcile_loop() -> None:
|
||||
logger.info("Discord stream reconciler started.")
|
||||
while True:
|
||||
await asyncio.sleep(settings.discord_reconcile_interval_seconds)
|
||||
try:
|
||||
await reconcile_once()
|
||||
except Exception as e:
|
||||
logger.error(f"Discord reconcile error: {e}")
|
||||
@@ -161,6 +161,26 @@ async def doc_delete(collection: str, doc_id: str) -> None:
|
||||
await asyncio.to_thread(ref.delete)
|
||||
|
||||
|
||||
def _run_transaction_sync(fn):
|
||||
"""Run ``fn(transaction)`` inside a Firestore transaction (retried on contention).
|
||||
|
||||
``fn`` must do every read (``ref.get(transaction=txn)`` /
|
||||
``query.stream(transaction=txn)``) before any write (``txn.set/update``) —
|
||||
Firestore's rule, and the reason this takes a function rather than being a
|
||||
context manager: the SDK re-runs the whole function when a read document
|
||||
changed underneath it. Bypasses the replay sandbox on purpose; nothing that
|
||||
uses it touches calls/incidents.
|
||||
"""
|
||||
txn = db.transaction()
|
||||
return fs.transactional(fn)(txn)
|
||||
|
||||
|
||||
async def run_transaction(fn):
|
||||
"""Async wrapper for _run_transaction_sync. Tests patch this with an
|
||||
in-memory fake (tests/fake_firestore.py)."""
|
||||
return await asyncio.to_thread(_run_transaction_sync, fn)
|
||||
|
||||
|
||||
async def doc_get_cached(collection: str, doc_id: str, ttl: float = 300.0) -> Optional[dict]:
|
||||
"""
|
||||
Like doc_get but backed by a short-lived in-memory TTL cache.
|
||||
|
||||
Reference in New Issue
Block a user