heystreamer/app/eventsub.py

import asyncio
import json
import logging
import os

import httpx
import websockets

from app.dispatch import dispatch_event

logger = logging.getLogger(__name__)

EVENTSUB_WS_URL = "wss://eventsub.wss.twitch.tv/ws"
TWITCH_CLIENT_ID = os.getenv("TWITCH_CLIENT_ID", "")
TWITCH_CLIENT_SECRET = os.getenv("TWITCH_CLIENT_SECRET", "")

SUBSCRIPTIONS = [
    ("channel.subscribe", "1", lambda bid: {"broadcaster_user_id": bid}),
    ("channel.subscription.message", "1", lambda bid: {"broadcaster_user_id": bid}),
    ("channel.subscription.gift", "1", lambda bid: {"broadcaster_user_id": bid}),
    ("channel.cheer", "1", lambda bid: {"broadcaster_user_id": bid}),
    ("channel.channel_points_custom_reward_redemption.add", "1", lambda bid: {"broadcaster_user_id": bid}),
    ("channel.follow", "2", lambda bid: {"broadcaster_user_id": bid, "moderator_user_id": bid}),
    ("channel.raid", "1", lambda bid: {"to_broadcaster_user_id": bid}),
]


async def run_eventsub(broadcaster_id: str, access_token: str) -> None:
    while True:
        try:
            async with websockets.connect(EVENTSUB_WS_URL) as ws:
                logger.info("EventSub connected")
                async for raw in ws:
                    msg = json.loads(raw)
                    msg_type = msg.get("metadata", {}).get("message_type")

                    if msg_type == "session_welcome":
                        session_id = msg["payload"]["session"]["id"]
                        token = await _subscribe_all(session_id, broadcaster_id, access_token)
                        if token:
                            access_token = token

                    elif msg_type == "notification":
                        await _handle_notification(msg["payload"])

                    elif msg_type == "session_reconnect":
                        reconnect_url = msg["payload"]["session"]["reconnect_url"]
                        logger.info("EventSub reconnect requested")
                        break

        except asyncio.CancelledError:
            logger.info("EventSub task cancelled")
            return
        except Exception as e:
            logger.warning(f"EventSub error: {e}, reconnecting in 10s")
            await asyncio.sleep(10)


async def _subscribe_all(session_id: str, broadcaster_id: str, access_token: str) -> str | None:
    transport = {"method": "websocket", "session_id": session_id}
    condition = {"broadcaster_user_id": broadcaster_id}
    headers = {
        "Authorization": f"Bearer {access_token}",
        "Client-Id": TWITCH_CLIENT_ID,
        "Content-Type": "application/json",
    }

    async with httpx.AsyncClient() as client:
        for sub_type, version, condition_fn in SUBSCRIPTIONS:
            condition = condition_fn(broadcaster_id)
            r = await client.post(
                "https://api.twitch.tv/helix/eventsub/subscriptions",
                headers=headers,
                json={"type": sub_type, "version": version, "condition": condition, "transport": transport},
            )
            if r.status_code == 401:
                new_token = await _refresh_token(access_token)
                if new_token:
                    headers["Authorization"] = f"Bearer {new_token}"
                    await client.post(
                        "https://api.twitch.tv/helix/eventsub/subscriptions",
                        headers=headers,
                        json={"type": sub_type, "version": version, "condition": condition, "transport": transport},
                    )
                    return new_token
            elif r.status_code not in (200, 202, 409):
                logger.warning(f"Failed to subscribe to {sub_type}: {r.status_code} {r.text}")
    return None


async def _refresh_token(old_token: str) -> str | None:
    from app.database import engine
    from app.models import User
    from sqlmodel import Session, select

    with Session(engine) as session:
        user = session.exec(select(User).where(User.access_token == old_token)).first()
        if not user:
            return None
        async with httpx.AsyncClient() as client:
            r = await client.post("https://id.twitch.tv/oauth2/token", data={
                "client_id": TWITCH_CLIENT_ID,
                "client_secret": TWITCH_CLIENT_SECRET,
                "grant_type": "refresh_token",
                "refresh_token": user.refresh_token,
            })
            if r.status_code != 200:
                return None
            data = r.json()
            user.access_token = data["access_token"]
            user.refresh_token = data.get("refresh_token", user.refresh_token)
            session.add(user)
            session.commit()
            return user.access_token


async def _handle_notification(payload: dict) -> None:
    sub_type = payload["subscription"]["type"]
    event = payload["event"]
    username = event.get("user_name") or event.get("user_login")

    if sub_type in ("channel.subscribe", "channel.subscription.message"):
        await dispatch_event("sub", username=username, event_detail=event.get("tier", "1000"))

    elif sub_type == "channel.subscription.gift":
        await dispatch_event("gift_sub", username=username, event_detail=event.get("tier", "1000"))

    elif sub_type == "channel.cheer":
        await dispatch_event("bits", username=username, bits_amount=event.get("bits", 0))

    elif sub_type == "channel.channel_points_custom_reward_redemption.add":
        reward_id = event.get("reward", {}).get("id")
        await dispatch_event("channel_points", username=username, event_detail=reward_id)

    elif sub_type == "channel.follow":
        await dispatch_event("follow", username=username)

    elif sub_type == "channel.raid":
        raider = event.get("from_broadcaster_user_name") or event.get("from_broadcaster_user_login")
        viewers = str(event.get("viewers", ""))
        await dispatch_event("raid", username=raider, event_detail=viewers)