diff --git a/telethon/client/telegrambaseclient.py b/telethon/client/telegrambaseclient.py index e980d811..d0608814 100644 --- a/telethon/client/telegrambaseclient.py +++ b/telethon/client/telegrambaseclient.py @@ -1,10 +1,10 @@ import abc import asyncio +import inspect import logging import platform import sys import time -import inspect from datetime import timedelta, datetime from .. import version @@ -209,7 +209,8 @@ class TelegramBaseClient(abc.ABC): retries=self._connection_retries, auto_reconnect=self._auto_reconnect, update_callback=self._handle_update, - auth_key_callback=self._auth_key_callback + auth_key_callback=self._auth_key_callback, + auto_reconnect_callback=self._handle_auto_reconnect ) # Cache :tl:`ExportedAuthorization` as ``dc_id: MTProtoState`` @@ -456,4 +457,8 @@ class TelegramBaseClient(abc.ABC): def _update_loop(self): raise NotImplementedError + @abc.abstractmethod + async def _handle_auto_reconnect(self): + raise NotImplementedError + # endregion diff --git a/telethon/client/updates.py b/telethon/client/updates.py index beaa5119..fe47e87d 100644 --- a/telethon/client/updates.py +++ b/telethon/client/updates.py @@ -1,12 +1,12 @@ import asyncio +import inspect import itertools import logging import random import time -import warnings from .users import UserMethods -from .. import events, utils +from .. import events, utils, errors from ..tl import types, functions __log__ = logging.getLogger(__name__) @@ -265,4 +265,14 @@ class UpdateMethods(UserMethods): __log__.exception('Unhandled exception on {}' .format(callback.__name__)) + async def _handle_auto_reconnect(self): + # Upon reconnection, we want to send getState + # for Telegram to keep sending us updates. + try: + __log__.info('Asking for the current state after reconnect...') + state = await self(functions.updates.GetStateRequest()) + __log__.info('Got new state! %s', state) + except errors.RPCError as e: + __log__.info('Failed to get current state: %r', e) + # endregion diff --git a/telethon/network/mtprotosender.py b/telethon/network/mtprotosender.py index fce18903..5eff89cf 100644 --- a/telethon/network/mtprotosender.py +++ b/telethon/network/mtprotosender.py @@ -42,7 +42,7 @@ class MTProtoSender: """ def __init__(self, state, connection, loop, *, retries=5, auto_reconnect=True, update_callback=None, - auth_key_callback=None): + auth_key_callback=None, auto_reconnect_callback=None): self.state = state self._connection = connection self._loop = loop @@ -52,6 +52,7 @@ class MTProtoSender: self._auto_reconnect = auto_reconnect self._update_callback = update_callback self._auth_key_callback = auth_key_callback + self._auto_reconnect_callback = auto_reconnect_callback # Whether the user has explicitly connected or disconnected. # @@ -307,6 +308,9 @@ class MTProtoSender: for m in self._pending_messages.values(): self._send_queue.put_nowait(m) + if self._auto_reconnect_callback: + self._loop.create_task(self._auto_reconnect_callback()) + break except ConnectionError: __log__.info('Failed reconnection retry %d/%d', retry, retries)