Update and persist session state on successful login

.gitignore

@@ -6,6 +6,7 @@
 
 # User session
 *.session
+sessions/
 /usermedia/
 
 # Builds and testing

telethon/_client/auth.py

@@ -253,7 +253,7 @@ async def sign_in(
         self._tos = result.terms_of_service
         raise errors.PhoneNumberUnoccupiedError(request=request)
 
-    return _on_login(self, result.user)
+    return await _update_session_state(self, result.user)
 
 async def sign_up(
         self: 'TelegramClient',
@@ -306,17 +306,28 @@ async def sign_up(
         await self(
             _tl.fn.help.AcceptTermsOfService(self._tos.id))
 
-    return _on_login(self, result.user)
+    return await _update_session_state(self, result.user)
 
-def _on_login(self, user):
+async def _update_session_state(self, user, save=True):
     """
     Callback called whenever the login or sign up process completes.
     Returns the input user parameter.
     """
-    self._bot = bool(user.bot)
-    self._self_input_peer = utils.get_input_peer(user, allow_self=False)
     self._authorized = True
 
+    self._session_state.user_id = user.id
+    self._session_state.bot = user.bot
+
+    state = await self(_tl.fn.updates.GetState())
+    self._session_state.pts = state.pts
+    self._session_state.qts = state.qts
+    self._session_state.date = int(state.date.timestamp())
+    self._session_state.seq = state.seq
+
+    await self.session.set_state(self._session_state)
+    if save:
+        await self.session.save()
+
     return user
 
 async def send_code_request(

telethon/_client/telegrambaseclient.py

@@ -385,16 +385,7 @@ async def connect(self: 'TelegramClient') -> None:
         # logged in). Attempt to fetch the user now. If it works, also get the update state.
         me = await self.get_me()
         if me:
-            self._session_state.user_id = me.id
-            self._session_state.bot = me.bot
-
-            state = await self(_tl.fn.updates.GetState())
-            self._session_state.pts = state.pts
-            self._session_state.qts = state.qts
-            self._session_state.date = int(state.date.timestamp())
-            self._session_state.seq = state.seq
-
-            await self.session.set_state(self._session_state)
+            await self._update_session_state(me, save=False)
 
     await self.session.save()
 

telethon/_client/telegramclient.py

@@ -3598,8 +3598,8 @@ class TelegramClient:
     async def _handle_auto_reconnect(self: 'TelegramClient'):
         return await updates._handle_auto_reconnect(**locals())
 
-    def _self_id(self: 'TelegramClient') -> typing.Optional[int]:
-        return users._self_id(**locals())
+    async def _update_session_state(self, user, save=True):
+        return await auth._update_session_state(**locals())
 
     # endregion Private