Send getState upon successful reconnection

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

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

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)