From 38857106893401106595d98b4558040b5519a83a Mon Sep 17 00:00:00 2001 From: Interitio Date: Sat, 1 Nov 2025 10:58:26 +1000 Subject: [PATCH] feat: Add websocket channel --- hyperfocus/channel.py | 107 +++++++++++++++++++++++++++++++++ hyperfocus/twitch/component.py | 52 ++++++++++++---- 2 files changed, 146 insertions(+), 13 deletions(-) create mode 100644 hyperfocus/channel.py diff --git a/hyperfocus/channel.py b/hyperfocus/channel.py new file mode 100644 index 0000000..0ac5dd6 --- /dev/null +++ b/hyperfocus/channel.py @@ -0,0 +1,107 @@ +from collections import defaultdict +from typing import Optional, TypeAlias, TypedDict +import json + +from datetime import datetime, timedelta + +from meta.sockets import Channel +from modules.profiles.profiles.profiles import ProfilesRegistry + +from . import logger +from .data import Hyperfocuser, HyperfocusData + +ISOTimestamp: TypeAlias = str + + +class HyperfocusedPayload(TypedDict): + userid: str + user_name: str + started_at: ISOTimestamp + ends_at: ISOTimestamp + + +async def prepare_hyperfocuser(profiler: ProfilesRegistry, hyperfocuser: Hyperfocuser): + profile = await profiler.get_profile(hyperfocuser.profileid) + assert profile is not None + + return HyperfocusedPayload( + userid=str(hyperfocuser.profileid), + user_name=profile.nickname or "Unknown", + started_at=hyperfocuser.started_at.isoformat(), + ends_at=hyperfocuser.ends_at.isoformat(), + ) + + +class FocusChannel(Channel): + name = "HyperFocus" + + def __init__(self, profiler: ProfilesRegistry, focusdata: HyperfocusData, **kwargs): + super().__init__(**kwargs) + + self.profiler = profiler + self.focusdata = focusdata + + # communityid -> listening websockets + self.communities = defaultdict(set) + + async def on_connection(self, websocket, event): + if not (cidstr := event.get("community")): + raise ValueError("Hyperfocus browser missing communityid") + elif not cidstr.isdigit(): + raise ValueError("Community id provided is not an integer") + cid = int(cidstr) + community = await self.profiler.get_community(cid) + if community is None: + raise ValueError("Unknown community provided") + + await super().on_connection(websocket, event) + self.communities[cid].add(websocket) + + focus_rows = await Hyperfocuser.fetch_where(started_in=cid) + await self.send_hyperfocus_put(cid, focus_rows) + + async def del_connection(self, websocket): + for wss in self.communities.values(): + wss.discard(websocket) + await super().del_connection(websocket) + + async def send_hyperfocus_patch( + self, communityid: int, focuser: Hyperfocuser, websocket=None + ): + for ws in (websocket,) if websocket else self.communities[communityid]: + await self.send_event( + { + "type": "DO", + "method": "patchFocus", + "args": await prepare_hyperfocuser(self.profiler, focuser), + }, + websocket=ws, + ) + + async def send_hyperfocus_del(self, profileid: int, websocket=None): + await self.send_event( + { + "type": "DO", + "method": "delFocus", + "args": {"userid": profileid}, + }, + websocket=websocket, + ) + + async def send_hyperfocus_put( + self, communityid: int, focusers: list[Hyperfocuser], websocket=None + ): + payload = [] + for focuser in focusers: + fpl = await prepare_hyperfocuser(self.profiler, focuser) + payload.append(fpl) + + for ws in (websocket,) if websocket else self.communities[communityid]: + await self.send_event( + {"type": "DO", "method": "putFocus", "args": payload}, + websocket=ws, + ) + + async def send_event(self, event, **kwargs): + logger.info(f"Sending websocket event: {json.dumps(event, indent=1)}") + await super().send_event(event, **kwargs) diff --git a/hyperfocus/twitch/component.py b/hyperfocus/twitch/component.py index f851ee6..2f1bd44 100644 --- a/hyperfocus/twitch/component.py +++ b/hyperfocus/twitch/component.py @@ -11,11 +11,13 @@ from twitchio.ext import commands as cmds from botdata import UserAuth from meta import Bot from meta.logger import log_wrap +from meta.sockets import register_channel from utils.lib import parse_dur, strfdelta, utc_now from . import logger from ..data import HyperfocusData, Hyperfocuser +from ..channel import FocusChannel # Default requested scopes for joining a channel @@ -33,6 +35,9 @@ class FocusComponent(cmds.Component): def __init__(self, bot: Bot): self.bot = bot self.data = bot.dbconn.load_registry(HyperfocusData()) + self.channel = FocusChannel(self.bot.profiles.profiles, self.data) + + register_channel(self.channel.name, self.channel) self._last_deleted: dict[int, datetime] = {} @@ -96,22 +101,41 @@ class FocusComponent(cmds.Component): ) try: await self.focus_delete_message(payload) - if notify: - await payload.broadcaster.send_message( - f"@{payload.chatter.name} Stay focused! " - "(You are in !hyperfocus, use !unfocus to come back if you need to!)", - sender=self.bot.bot_id, - ) + deleted = True except Exception: logger.warning( f"Failed to delete a hyperfocus message: {payload!r}", exc_info=True ) - if notify: - await payload.broadcaster.send_message( - f"@{payload.chatter.name} Stay focused! ", - sender=self.bot.bot_id, + deleted = False + + if notify: + self._last_deleted[profile.profileid] = utc_now() + try: + if deleted: + await payload.broadcaster.send_message( + f"@{payload.chatter.name} Stay focused! " + "(You are in !hyperfocus, use !unfocus to come back if you need to!)", + sender=self.bot.bot_id, + ) + else: + await payload.broadcaster.send_message( + f"@{payload.chatter.name} Stay focused! ", + sender=self.bot.bot_id, + ) + except twitchio.exceptions.HTTPException: + logger.warning( + f"Failed to notify user of hyperfocus deletion: {payload!r}", + exc_info=True, ) - self._last_deleted[profile.profileid] = utc_now() + + if hyperfocused: + # Send an update to the channel + # TODO: Nicer and more efficient caching of which channel has which users + # TODO: Possible race condition with the commands. Should use locks. + comm = await self.bot.profiles.fetch_community( + payload.broadcaster, touch=True + ) + await self.channel.send_hyperfocus_patch(comm.communityid, hyperfocused) # ------ Commands ----- @cmds.command( @@ -150,14 +174,15 @@ class FocusComponent(cmds.Component): comm = await self.bot.profiles.fetch_community(ctx.broadcaster, touch=True) await Hyperfocuser.table.delete_where(profileid=pid) - await Hyperfocuser.create( + focuser = await Hyperfocuser.create( profileid=pid, started_at=now, ends_at=end_at, started_in=comm.communityid, ) - # TODO: Update channel + await self.channel.send_hyperfocus_patch(comm.communityid, focuser) + minutes = ceil(dur / 60) await ctx.reply( f"{ctx.chatter.name} has gone into HYPERFOCUS! " @@ -171,6 +196,7 @@ class FocusComponent(cmds.Component): row = await Hyperfocuser.fetch(profile.profileid) if row: await row.delete() + await self.channel.send_hyperfocus_del(profile.profileid) await ctx.reply( "Welcome back from focus, hope it went well!"