Add pixels discord stream

Author: jb3

bot/constants.py

@@ -98,6 +98,7 @@ class _Channels(EnvConfig, env_prefix="CHANNEL_"):
     voice_chat_0: int = 412357430186344448
     voice_chat_1: int = 799647045886541885
     roles: int = 851270062434156586
+    pixels_stream: int = 1541576692986941563
 
 
 Channels = _Channels()

bot/exts/pixels.py

@@ -1,11 +1,17 @@
+import asyncio
+import json
 import re
 import textwrap
+from os import environ
 from pathlib import Path
 
-from discord import Message
+import discord
+from discord import Message, ui
 from discord.ext import commands
+from pydis_core.utils import scheduling
 from pydis_core.utils.logging import get_logger
 
+from bot import constants
 from bot.bot import SirRobin
 
 log = get_logger(__name__)
@@ -20,12 +26,38 @@
 """)
 
 PIXELS_API_TOKEN_PATH = Path("/var/run/secrets/pixels-api/token")
+PIXELS_URL = environ.get("PIXELS_URL", "https://pixels.pythondiscord.com")
+
+class PixelsStreamView(ui.LayoutView):
+    """LayoutView containing placement containers."""
+
+    def __init__(self, containers: list[ui.Container]) -> None:
+        super().__init__()
+        for container in containers:
+            self.add_item(container)
+
 
 class Pixels(commands.Cog):
     """Utilities for working with Pixels API."""
 
     def __init__(self, bot: SirRobin):
         self.bot = bot
+        self.stream_task: asyncio.Task | None = None
+        self.flush_task: asyncio.Task | None = None
+        self.bucket: list[ui.Container] = []
+        self.bucket_lock = asyncio.Lock()
+
+    async def cog_load(self) -> None:
+        """Start background tasks for streaming and flushing placement buckets."""
+        self.stream_task = scheduling.create_task(self.consume_stream())
+        self.flush_task = scheduling.create_task(self.flush_bucket_loop())
+
+    async def cog_unload(self) -> None:
+        """Cancel background tasks on cog unload."""
+        if self.stream_task:
+            self.stream_task.cancel()
+        if self.flush_task:
+            self.flush_task.cancel()
 
     def get_request_headers(self) -> dict[str, str]:
         """Get the Pixels API request headers from the projected service account file."""
@@ -34,6 +66,130 @@ def get_request_headers(self) -> dict[str, str]:
 
         return {"Authorization": f"Token {token}"}
 
+    async def consume_stream(self) -> None:
+        """Consume the Pixels API event stream."""
+        await self.bot.wait_until_ready()
+        stream_url = f"{PIXELS_URL.rstrip('/')}/api/stream?no_current_state=true"
+
+        while True:
+            try:
+                headers = self.get_request_headers()
+                async with self.bot.http_session.get(stream_url, headers=headers) as resp:
+                    if resp.status != 200:
+                        log.error(f"Failed to connect to Pixels stream. Status code: {resp.status}")
+                        await asyncio.sleep(5)
+                        continue
+
+                    log.info("Connected to Pixels API stream.")
+                    current_event: str | None = None
+                    current_data: list[str] = []
+
+                    async for line_bytes in resp.content:
+                        line = line_bytes.decode("utf-8").rstrip("\r\n")
+                        if not line:
+                            if current_event and current_data:
+                                data_str = "\n".join(current_data)
+                                await self.handle_stream_event(current_event, data_str)
+                            current_event = None
+                            current_data = []
+                        elif line.startswith("event:"):
+                            current_event = line[6:].strip()
+                        elif line.startswith("data:"):
+                            current_data.append(line[5:].strip())
+            except asyncio.CancelledError:
+                log.info("Pixels stream consumer task cancelled.")
+                break
+            except Exception:
+                log.exception("Error consuming Pixels API stream.")
+                await asyncio.sleep(5)
+
+    async def handle_stream_event(self, event_type: str, data_str: str) -> None:
+        """Handle an incoming event from the Pixels API stream."""
+        try:
+            payload = json.loads(data_str)
+        except json.JSONDecodeError:
+            log.warning(f"Failed to decode stream event data for event {event_type}: {data_str}")
+            return
+
+        user_id = payload.get("user_id")
+        rgb = payload.get("rgb")
+        x = payload.get("x")
+        y = payload.get("y")
+        if user_id is None or not rgb or x is None or y is None:
+            log.warning(f"Event payload missing required fields: {payload}")
+            return
+
+        user_id_int = int(user_id)
+        user = self.bot.get_user(user_id_int)
+        if user is None:
+            try:
+                user = await self.bot.fetch_user(user_id_int)
+            except discord.HTTPException:
+                user = None
+
+        if user:
+            user_name = user.name
+            avatar_url = user.display_avatar.url
+        else:
+            user_name = "Unknown User"
+            avatar_url = self.bot.user.display_avatar.url
+
+        rgb_hex = str(rgb).lstrip("#")
+        try:
+            colour = discord.Colour(int(rgb_hex, 16))
+        except ValueError:
+            colour = discord.Colour.default()
+
+        container = ui.Container(accent_colour=colour)
+        section = ui.Section(
+            f"{user_name} (`{user_id_int}`) at ({x}, {y})",
+            accessory=ui.Thumbnail(avatar_url),
+        )
+        container.add_item(section)
+
+        await self.add_container_to_bucket(container)
+
+    async def add_container_to_bucket(self, container: ui.Container) -> None:
+        """Add a container to the batch bucket. If 10 containers are reached, send the bucket immediately."""
+        async with self.bucket_lock:
+            self.bucket.append(container)
+            if len(self.bucket) >= 10:
+                await self._send_bucket_locked()
+
+    async def flush_bucket_loop(self) -> None:
+        """Flush the bucket every 10 seconds if it contains any containers."""
+        await self.bot.wait_until_ready()
+        while True:
+            try:
+                await asyncio.sleep(10)
+                async with self.bucket_lock:
+                    if self.bucket:
+                        await self._send_bucket_locked()
+            except asyncio.CancelledError:
+                break
+            except Exception:
+                log.exception("Error in Pixels flush bucket loop.")
+
+    async def _send_bucket_locked(self) -> None:
+        """Send containers in the bucket as components v2 LayoutViews in batches of up to 10."""
+        while self.bucket:
+            containers_to_send = self.bucket[:10]
+            self.bucket = self.bucket[10:]
+
+            channel = self.bot.get_channel(constants.Channels.pixels_stream)
+            if channel is None:
+                try:
+                    channel = await self.bot.fetch_channel(constants.Channels.pixels_stream)
+                except discord.HTTPException:
+                    log.error(f"Could not fetch pixels stream channel (ID: {constants.Channels.pixels_stream}).")
+                    return
+
+            view = PixelsStreamView(containers_to_send)
+            try:
+                await channel.send(view=view)
+            except discord.HTTPException:
+                log.exception(f"Failed to post Pixels placement bucket to channel {channel.id}.")
+
     @commands.Cog.listener()
     async def on_message(self, message: Message) -> None:
         """Check for exposed Pixels API keys and JWTs in messages."""
@@ -51,7 +207,7 @@ async def on_message(self, message: Message) -> None:
 
             for jwt in found_jwts:
                 async with self.bot.http_session.post(
-                    "https://pixels.pythondiscord.com/api/invalidate_credential",
+                    f"{PIXELS_URL.rstrip('/')}/api/invalidate_credential",
                     headers=headers,
                     json={"jwt": jwt},
                 ) as resp:
@@ -67,7 +223,7 @@ async def on_message(self, message: Message) -> None:
 
             for api_key in found_api_keys:
                 async with self.bot.http_session.post(
-                    "https://pixels.pythondiscord.com/api/invalidate_credential",
+                    f"{PIXELS_URL.rstrip('/')}/api/invalidate_credential",
                     headers=headers,
                     json={"api_key": api_key},
                 ) as resp:

tests/test_pixels.py

@@ -0,0 +1,54 @@
+import json
+import unittest
+from unittest.mock import MagicMock
+
+import discord
+
+from bot.exts.pixels import Pixels, PixelsStreamView
+from tests.helpers import MockBot, MockTextChannel
+
+
+class PixelsCogTests(unittest.IsolatedAsyncioTestCase):
+    """Tests for Pixels cog streaming and batching functionality."""
+
+    def setUp(self) -> None:
+        self.bot = MockBot()
+        self.cog = Pixels(self.bot)
+
+    async def test_handle_stream_event_creates_container(self) -> None:
+        """Placement events should create a container and add it to bucket."""
+        user = MagicMock(spec=discord.User)
+        user.name = "TestUser"
+        user.id = 123456789
+        user.display_avatar.url = "http://avatar.url/1.png"
+        self.bot.get_user.return_value = user
+
+        event_data = json.dumps({"x": 10, "y": 20, "rgb": "FF0000", "user_id": 123456789})
+        await self.cog.handle_stream_event("UPDATE", event_data)
+
+        self.assertEqual(len(self.cog.bucket), 1)
+        container = self.cog.bucket[0]
+        self.assertIsInstance(container, discord.ui.Container)
+        self.assertEqual(container.accent_colour.value, 0xFF0000)
+
+    async def test_bucket_sends_when_10_containers_reached(self) -> None:
+        """Bucket should flush automatically when 10 containers are reached."""
+        channel = MockTextChannel()
+        self.bot.get_channel.return_value = channel
+
+        user = MagicMock(spec=discord.User)
+        user.name = "TestUser"
+        user.id = 123456789
+        user.display_avatar.url = "http://avatar.url/1.png"
+        self.bot.get_user.return_value = user
+
+        event_data = json.dumps({"x": 10, "y": 20, "rgb": "00FF00", "user_id": 123456789})
+        for _ in range(10):
+            await self.cog.handle_stream_event("UPDATE", event_data)
+
+        self.assertEqual(len(self.cog.bucket), 0)
+        channel.send.assert_called_once()
+        _, kwargs = channel.send.call_args
+        self.assertIn("view", kwargs)
+        self.assertIsInstance(kwargs["view"], PixelsStreamView)
+        self.assertEqual(len(kwargs["view"].children), 10)