diff --git a/misskey/input.py b/misskey/input.py index 14db900..7daeed1 100644 --- a/misskey/input.py +++ b/misskey/input.py @@ -242,5 +242,7 @@ class MisskeyInputService(MisskeyService, InputService): streaming: str = f"{'wss' if self.url.startswith('https') else 'ws'}://{self.url.split('://', 1)[1]}" url: str = f"{streaming}/streaming?i={self.options.token}" - async for msg in listen_websocket(url, self.log): + async for msg in listen_websocket( + url, self.log, on_connect=self._subscribe_to_home + ): self.submitter(lambda: self._accept_msg(msg)) diff --git a/util/util.py b/util/util.py index 04b6946..9dc8513 100644 --- a/util/util.py +++ b/util/util.py @@ -3,7 +3,7 @@ import logging import os import sys from collections.abc import AsyncIterator, Callable -from typing import Any, cast +from typing import Any, Awaitable, cast from urllib.parse import urlparse import websockets @@ -51,6 +51,7 @@ async def listen_websocket( url: str, log: logging.Logger, *, + on_connect: Callable[[websockets.ClientConnection], Awaitable[None]] | None = None, headers: dict[str, str] | None = None, ping_interval: int = 20, ping_timeout: int = 10, @@ -71,6 +72,8 @@ async def listen_websocket( ping_timeout=ping_timeout, close_timeout=close_timeout, ) as ws: + if on_connect: + await on_connect(ws) log.info("Listening to '%s'...", label) backoff = 1.0 async for msg in ws: