Improve session storage lifecycle

client/src/telethon/_impl/client/client/auth.py

@@ -260,4 +260,9 @@ async def check_password(
 async def sign_out(self: Client) -> None:
     await self(functions.auth.log_out())
 
-    await self._storage.delete()
+    self._chat_hashes.clear()
+    self._message_box.reset()
+
+    self._session.user = None
+    self._session.state = None
+    await self._storage.save(self._session)

client/src/telethon/_impl/client/client/net.py

@@ -230,8 +230,13 @@ async def disconnect(self: Client) -> None:
             "unhandled exception during disconnect; this is a bug"
         )
 
+    try:
+        if self._session.user:
+            # Only save if we haven't logged out (prevents double-save)
             self._session.state = self._message_box.session_state()
             await self._storage.save(self._session)
+    finally:
+        await self._storage.close()
 
 
 async def invoke_request(

client/src/telethon/_impl/session/chat/hash_cache.py

@@ -33,6 +33,11 @@ class ChatHashCache:
         else:
             return None
 
+    def clear(self) -> None:
+        self._hash_map.clear()
+        self._self_id = None
+        self._self_bot = False
+
     def _has(self, id: int) -> bool:
         return id in self._hash_map
 

client/src/telethon/_impl/session/message_box/messagebox.py

@@ -103,6 +103,14 @@ class MessageBox:
         self.getting_diff_for.clear()
         self.next_deadline = ENTRY_ACCOUNT
 
+    def reset(self) -> None:
+        self.map.clear()
+        self.date = epoch()
+        self.seq = NO_SEQ
+        self.possible_gaps.clear()
+        self.getting_diff_for.clear()
+        self.next_deadline = None
+
     def session_state(self) -> UpdateState:
         return UpdateState(
             pts=self.map[ENTRY_ACCOUNT].pts if ENTRY_ACCOUNT in self.map else NO_PTS,

client/src/telethon/_impl/session/storage/memory.py

@@ -24,5 +24,5 @@ class MemorySession(Storage):
     async def save(self, session: Session) -> None:
         self.session = session
 
-    async def delete(self) -> None:
-        self.session = None
+    async def close(self) -> None:
+        pass

client/src/telethon/_impl/session/storage/sqlite.py

@@ -52,14 +52,11 @@ class SqliteSession(Storage):
         conn = self._current_conn()
         with conn:
             self._save_v10(conn.cursor(), session)
-        conn.close()
-        self._conn = None
 
-    async def delete(self) -> None:
+    async def close(self) -> None:
         if self._conn:
             self._conn.close()
             self._conn = None
-        self._path.unlink()
 
     def _current_conn(self) -> sqlite3.Connection:
         if self._conn is None:

client/src/telethon/_impl/session/storage/storage.py

@@ -14,7 +14,7 @@ class Storage(abc.ABC):
         """
         Load the :class:`Session` instance, if any.
 
-        This method is called by the library prior to `connect`.
+        This method is called by the library prior to ``connect``.
 
         :return: The previously-saved session.
         """
@@ -24,18 +24,21 @@ class Storage(abc.ABC):
         """
         Save the :class:`Session` instance to persistent storage.
 
-        This method is called by the library post `disconnect`.
+        This method is called by the library after significant changes to the session,
+        such as login, logout, or to persist the update state prior to disconnection.
 
         :param session:
             The session information that should be persisted.
         """
 
     @abc.abstractmethod
-    async def delete(self) -> None:
+    async def close(self) -> None:
         """
-        Delete the saved `Session`.
+        Close the :class:`Session` instance, if it was still open.
 
-        This method is called by the library post `log_out`.
+        This method is called by the library post ``disconnect``,
+        even if the call to :meth:`save` failed.
 
-        Note that both :meth:`load` and :meth:`save` may still be called after.
+        Note that :meth:`load` may still be called after,
+        in which case the session should be reopened.
         """