From d6d5ac10e66f0e80b8d4215f10d552545b7ba514 Mon Sep 17 00:00:00 2001 From: cottongin Date: Thu, 12 Mar 2026 07:51:55 -0400 Subject: [PATCH] fix: separate bot vs viewer WebSocket connections, add client identification The dashboard's own WS connection was being counted as a bot subscriber, causing "1 bot connected" with no bots actually present. Now WS clients send a role ("bot" or "viewer") in the subscribe message. Only bots count toward the subscriber total. Bot plugins also send a configurable client_id so the dashboard shows which specific bots are connected. Made-with: Cursor --- plugins/limnoria/NtrPlaylist/config.py | 9 ++++ plugins/limnoria/NtrPlaylist/plugin.py | 8 ++- plugins/sopel/ntr_playlist.py | 9 +++- src/ntr_fetcher/dashboard.py | 10 +++- src/ntr_fetcher/static/dashboard.html | 26 +++++++-- src/ntr_fetcher/websocket.py | 67 +++++++++++++++++------ tests/test_dashboard.py | 15 +++++- tests/test_websocket.py | 75 ++++++++++++++++++++++---- 8 files changed, 184 insertions(+), 35 deletions(-) diff --git a/plugins/limnoria/NtrPlaylist/config.py b/plugins/limnoria/NtrPlaylist/config.py index 59c3766..d95304a 100644 --- a/plugins/limnoria/NtrPlaylist/config.py +++ b/plugins/limnoria/NtrPlaylist/config.py @@ -61,3 +61,12 @@ conf.registerGlobalValue( """IRC channel to send announce messages to.""", ), ) + +conf.registerGlobalValue( + NtrPlaylist, + "clientId", + registry.String( + "limnoria", + """Identifier for this bot when connecting to the announce WebSocket.""", + ), +) diff --git a/plugins/limnoria/NtrPlaylist/plugin.py b/plugins/limnoria/NtrPlaylist/plugin.py index 6240cf1..f32a93c 100644 --- a/plugins/limnoria/NtrPlaylist/plugin.py +++ b/plugins/limnoria/NtrPlaylist/plugin.py @@ -144,6 +144,7 @@ class NtrPlaylist(callbacks.Plugin): ws_url = self.registryValue("wsUrl") token = self.registryValue("adminToken") channel = self.registryValue("announceChannel") + client_id = self.registryValue("clientId") or "limnoria" if not ws_url or not token: LOGGER.warning("wsUrl or adminToken not configured, WS listener sleeping") @@ -154,7 +155,12 @@ class NtrPlaylist(callbacks.Plugin): try: ws = websocket.WebSocket() ws.connect(ws_url, timeout=10) - ws.send(json.dumps({"type": "subscribe", "token": token})) + ws.send(json.dumps({ + "type": "subscribe", + "token": token, + "role": "bot", + "client_id": client_id, + })) LOGGER.info("Connected to announce WebSocket at %s", ws_url) backoff = 5 diff --git a/plugins/sopel/ntr_playlist.py b/plugins/sopel/ntr_playlist.py index 0898d81..368d66b 100644 --- a/plugins/sopel/ntr_playlist.py +++ b/plugins/sopel/ntr_playlist.py @@ -24,6 +24,7 @@ class NtrPlaylistSection(types.StaticSection): display_timezone = types.ValidatedAttribute("display_timezone", default="America/New_York") ws_url = types.ValidatedAttribute("ws_url", default="ws://127.0.0.1:8000/ws/announce") announce_channel = types.ValidatedAttribute("announce_channel", default="#sewerchat") + client_id = types.ValidatedAttribute("client_id", default="sopel") _ws_stop = None @@ -54,6 +55,7 @@ def _ws_listener(bot): ws_url = bot.settings.ntr_playlist.ws_url token = bot.settings.ntr_playlist.admin_token channel = bot.settings.ntr_playlist.announce_channel + client_id = bot.settings.ntr_playlist.client_id or "sopel" if not ws_url or not token: LOGGER.warning("ws_url or admin_token not configured, WS listener sleeping") @@ -64,7 +66,12 @@ def _ws_listener(bot): try: ws = websocket.WebSocket() ws.connect(ws_url, timeout=10) - ws.send(json.dumps({"type": "subscribe", "token": token})) + ws.send(json.dumps({ + "type": "subscribe", + "token": token, + "role": "bot", + "client_id": client_id, + })) LOGGER.info("Connected to announce WebSocket at %s", ws_url) backoff = 5 diff --git a/src/ntr_fetcher/dashboard.py b/src/ntr_fetcher/dashboard.py index 3cc3bbd..fc6e2f7 100644 --- a/src/ntr_fetcher/dashboard.py +++ b/src/ntr_fetcher/dashboard.py @@ -128,7 +128,13 @@ def create_dashboard_router( await websocket.close(code=4001, reason="Invalid token") return - manager.add_subscriber(websocket) + role = auth_msg.get("role", "bot") + client_id = auth_msg.get("client_id", "") + remote_addr = "" + if websocket.client: + remote_addr = websocket.client.host or "" + + manager.add_client(websocket, role=role, client_id=client_id, remote_addr=remote_addr) await manager.broadcast_status() try: while True: @@ -138,7 +144,7 @@ def create_dashboard_router( except WebSocketDisconnect: pass finally: - manager.remove_subscriber(websocket) + manager.remove_client(websocket) try: await manager.broadcast_status() except Exception: diff --git a/src/ntr_fetcher/static/dashboard.html b/src/ntr_fetcher/static/dashboard.html index 11bc935..5121e6d 100644 --- a/src/ntr_fetcher/static/dashboard.html +++ b/src/ntr_fetcher/static/dashboard.html @@ -33,6 +33,11 @@ .toast.error { background: #e53935; } details summary h3 { display: inline; cursor: pointer; } .track-num { width: 40px; text-align: center; } + .client-tag { + display: inline-block; padding: 2px 8px; margin: 2px 4px; + border-radius: 4px; background: #2a2a2a; font-size: 0.8rem; + } + #client-detail { display: none; margin-top: 4px; } @@ -42,6 +47,7 @@ NtR Playlist Dashboard (connecting...) +
Logout @@ -161,16 +167,28 @@ } } - function updateStatus(count) { + function updateStatus(count, clients) { subscriberCount = count; const dot = document.getElementById("status-dot"); const sub = document.getElementById("sub-count"); + const detail = document.getElementById("client-detail"); if (count > 0) { dot.className = "status-dot connected"; sub.textContent = `(${count} bot${count > 1 ? "s" : ""} connected)`; + if (clients && clients.length > 0) { + detail.innerHTML = clients.map(c => { + const name = esc(c.client_id || "unknown"); + const addr = esc(c.remote_addr || "?"); + return `${name} (${addr})`; + }).join(" "); + detail.style.display = "block"; + } else { + detail.style.display = "none"; + } } else { dot.className = "status-dot disconnected"; sub.textContent = "(no bots connected)"; + detail.style.display = "none"; } document.querySelectorAll(".announce-btn").forEach(btn => { if (count === 0) { @@ -188,19 +206,19 @@ const proto = location.protocol === "https:" ? "wss:" : "ws:"; const ws = new WebSocket(`${proto}//${location.host}/ws/announce`); ws.onopen = () => { - ws.send(JSON.stringify({type: "subscribe", token: WS_TOKEN})); + ws.send(JSON.stringify({type: "subscribe", token: WS_TOKEN, role: "viewer"})); wsBackoff = 1000; }; ws.onmessage = (e) => { try { const data = JSON.parse(e.data); if (data.type === "status") { - updateStatus(data.subscribers); + updateStatus(data.subscribers, data.clients || []); } } catch {} }; ws.onclose = () => { - updateStatus(0); + updateStatus(0, []); document.getElementById("sub-count").textContent = "(reconnecting...)"; setTimeout(connectWS, wsBackoff); wsBackoff = Math.min(wsBackoff * 2, 60000); diff --git a/src/ntr_fetcher/websocket.py b/src/ntr_fetcher/websocket.py index 23df44e..5002d88 100644 --- a/src/ntr_fetcher/websocket.py +++ b/src/ntr_fetcher/websocket.py @@ -1,37 +1,74 @@ import logging +from dataclasses import dataclass, field +from datetime import datetime, timezone logger = logging.getLogger(__name__) +@dataclass +class Client: + websocket: object + role: str + client_id: str = "" + remote_addr: str = "" + connected_at: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) + + class AnnounceManager: def __init__(self): - self._subscribers: list = [] + self._clients: list[Client] = [] @property - def subscriber_count(self) -> int: - return len(self._subscribers) + def bot_count(self) -> int: + return sum(1 for c in self._clients if c.role == "bot") - def add_subscriber(self, websocket) -> None: - self._subscribers.append(websocket) - logger.info("Subscriber connected (%d total)", self.subscriber_count) + @property + def bot_clients(self) -> list[dict]: + return [ + { + "client_id": c.client_id, + "remote_addr": c.remote_addr, + "connected_at": c.connected_at, + } + for c in self._clients + if c.role == "bot" + ] - def remove_subscriber(self, websocket) -> None: - self._subscribers = [ws for ws in self._subscribers if ws is not websocket] - logger.info("Subscriber disconnected (%d total)", self.subscriber_count) + def add_client(self, websocket, role: str, client_id: str = "", remote_addr: str = "") -> None: + self._clients.append(Client( + websocket=websocket, + role=role, + client_id=client_id, + remote_addr=remote_addr, + )) + logger.info( + "Client connected: role=%s client_id=%s addr=%s (%d bots, %d total)", + role, client_id, remote_addr, self.bot_count, len(self._clients), + ) + + def remove_client(self, websocket) -> None: + removed = [c for c in self._clients if c.websocket is websocket] + self._clients = [c for c in self._clients if c.websocket is not websocket] + for c in removed: + logger.info( + "Client disconnected: role=%s client_id=%s (%d bots, %d total)", + c.role, c.client_id, self.bot_count, len(self._clients), + ) async def broadcast(self, message: dict) -> None: dead = [] - for ws in self._subscribers: + for client in self._clients: try: - await ws.send_json(message) + await client.websocket.send_json(message) except Exception: - dead.append(ws) - logger.warning("Removing dead subscriber") + dead.append(client.websocket) + logger.warning("Removing dead client: %s", client.client_id or client.role) for ws in dead: - self.remove_subscriber(ws) + self.remove_client(ws) async def broadcast_status(self) -> None: await self.broadcast({ "type": "status", - "subscribers": self.subscriber_count, + "subscribers": self.bot_count, + "clients": self.bot_clients, }) diff --git a/tests/test_dashboard.py b/tests/test_dashboard.py index 9a999d5..6cd91c2 100644 --- a/tests/test_dashboard.py +++ b/tests/test_dashboard.py @@ -162,12 +162,23 @@ def test_announce_invalid_position(client, db): # --- WebSocket tests --- -def test_ws_subscribe_with_valid_token(app): +def test_ws_subscribe_bot_with_valid_token(app): with TestClient(app) as c: with c.websocket_connect("/ws/announce") as ws: - ws.send_json({"type": "subscribe", "token": "test-token"}) + ws.send_json({"type": "subscribe", "token": "test-token", "role": "bot", "client_id": "test-bot"}) data = ws.receive_json() assert data["type"] == "status" + assert data["subscribers"] == 1 + assert data["clients"][0]["client_id"] == "test-bot" + + +def test_ws_subscribe_viewer_not_counted(app): + with TestClient(app) as c: + with c.websocket_connect("/ws/announce") as ws: + ws.send_json({"type": "subscribe", "token": "test-token", "role": "viewer"}) + data = ws.receive_json() + assert data["type"] == "status" + assert data["subscribers"] == 0 def test_ws_subscribe_with_invalid_token(app): diff --git a/tests/test_websocket.py b/tests/test_websocket.py index aa86d08..2fbefdb 100644 --- a/tests/test_websocket.py +++ b/tests/test_websocket.py @@ -7,12 +7,12 @@ def manager(): return AnnounceManager() -def test_no_subscribers_initially(manager): - assert manager.subscriber_count == 0 +def test_no_bots_initially(manager): + assert manager.bot_count == 0 @pytest.mark.asyncio -async def test_subscribe_and_broadcast(manager): +async def test_bot_subscribe_and_broadcast(manager): received = [] class FakeWS: @@ -20,15 +20,34 @@ async def test_subscribe_and_broadcast(manager): received.append(data) ws = FakeWS() - manager.add_subscriber(ws) - assert manager.subscriber_count == 1 + manager.add_client(ws, role="bot", client_id="test-bot", remote_addr="127.0.0.1") + assert manager.bot_count == 1 await manager.broadcast({"type": "announce", "message": "Now Playing: Song #1"}) assert len(received) == 1 assert received[0]["message"] == "Now Playing: Song #1" - manager.remove_subscriber(ws) - assert manager.subscriber_count == 0 + manager.remove_client(ws) + assert manager.bot_count == 0 + + +@pytest.mark.asyncio +async def test_viewer_not_counted_as_bot(manager): + received = [] + + class FakeWS: + async def send_json(self, data): + received.append(data) + + ws = FakeWS() + manager.add_client(ws, role="viewer") + assert manager.bot_count == 0 + + await manager.broadcast({"type": "announce", "message": "test"}) + assert len(received) == 1 + + manager.remove_client(ws) + assert manager.bot_count == 0 @pytest.mark.asyncio @@ -38,8 +57,44 @@ async def test_broadcast_skips_dead_connections(manager): raise Exception("connection closed") ws = DeadWS() - manager.add_subscriber(ws) - assert manager.subscriber_count == 1 + manager.add_client(ws, role="bot", client_id="dead-bot") + assert manager.bot_count == 1 await manager.broadcast({"type": "announce", "message": "test"}) - assert manager.subscriber_count == 0 + assert manager.bot_count == 0 + + +@pytest.mark.asyncio +async def test_bot_clients_returns_metadata(manager): + class FakeWS: + async def send_json(self, data): + pass + + ws = FakeWS() + manager.add_client(ws, role="bot", client_id="limnoria-prod", remote_addr="10.0.0.5") + clients = manager.bot_clients + assert len(clients) == 1 + assert clients[0]["client_id"] == "limnoria-prod" + assert clients[0]["remote_addr"] == "10.0.0.5" + assert "connected_at" in clients[0] + + +@pytest.mark.asyncio +async def test_status_broadcast_includes_clients(manager): + received = [] + + class FakeWS: + async def send_json(self, data): + received.append(data) + + bot = FakeWS() + viewer = FakeWS() + manager.add_client(bot, role="bot", client_id="my-bot", remote_addr="1.2.3.4") + manager.add_client(viewer, role="viewer") + + await manager.broadcast_status() + for msg in received: + assert msg["type"] == "status" + assert msg["subscribers"] == 1 + assert len(msg["clients"]) == 1 + assert msg["clients"][0]["client_id"] == "my-bot"