From 375e97d9636572c1f72c6f62b506b9817b9702b6 Mon Sep 17 00:00:00 2001
From: wulan17 <wulan17@komodos.id>
Date: Sat, 7 Jun 2025 20:58:36 +0700
Subject: [PATCH] Reapply "fix: handle connection closure and retry logic in
 session management"

This reverts commit 2c3fb1caa6877a31621d0535866672a13334d0ee.

Signed-off-by: wulan17 <wulan17@komodos.id>
---
 pyrogram/connection/transport/tcp/tcp.py | 17 +++++++++++++++++
 pyrogram/session/session.py              | 23 +++++++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/pyrogram/connection/transport/tcp/tcp.py b/pyrogram/connection/transport/tcp/tcp.py
index 13e9f806..a872c608 100644
--- a/pyrogram/connection/transport/tcp/tcp.py
+++ b/pyrogram/connection/transport/tcp/tcp.py
@@ -59,6 +59,13 @@ class TCP:
             self.loop = loop
         else:
             self.loop = asyncio.get_event_loop()
+        self._closed = True
+
+    @property
+    def closed(self) -> bool:
+        return (
+            self._closed or self.writer is None or self.writer.is_closing() or self.reader is None
+        )
 
     async def _connect_via_proxy(
         self,
@@ -126,11 +133,14 @@ class TCP:
     async def connect(self, address: Tuple[str, int]) -> None:
         try:
             await asyncio.wait_for(self._connect(address), TCP.TIMEOUT)
+            self._closed = False
         except asyncio.TimeoutError:  # Re-raise as TimeoutError. asyncio.TimeoutError is deprecated in 3.11
+            self._closed = True
             raise TimeoutError("Connection timed out")
 
     async def close(self) -> None:
         if self.writer is None:
+            self._closed = True
             return None
 
         try:
@@ -140,6 +150,7 @@ class TCP:
             log.info("Close exception: %s %s", type(e).__name__, e)
         finally:
             self.writer = None
+            self._closed = True
 
     async def send(self, data: bytes) -> None:
         async with self.lock:
@@ -151,9 +162,13 @@ class TCP:
                 await self.writer.drain()
             except Exception as e:
                 log.info("Send exception: %s %s", type(e).__name__, e)
+                self._closed = True
                 raise OSError(e)
 
     async def recv(self, length: int = 0) -> Optional[bytes]:
+        if self._closed or self.reader is None:
+            return None
+
         data = b""
 
         while len(data) < length:
@@ -163,11 +178,13 @@ class TCP:
                     TCP.TIMEOUT
                 )
             except (OSError, asyncio.TimeoutError):
+                self._closed = True
                 return None
             else:
                 if chunk:
                     data += chunk
                 else:
+                    self._closed = True
                     return None
 
         return data
diff --git a/pyrogram/session/session.py b/pyrogram/session/session.py
index 3424b42e..098d0cd9 100644
--- a/pyrogram/session/session.py
+++ b/pyrogram/session/session.py
@@ -432,6 +432,19 @@ class Session:
 
         while retries > 0:
             try:
+                if (
+                    self.connection is None
+                    or self.connection.protocol is None
+                    or getattr(self.connection.protocol, "closed", True)
+                ):
+                    log.warning(
+                        "[%s] Connection is closed or not established. Attempting to reconnect...",
+                        self.client.name,
+                    )
+                    await self.restart()
+                    await asyncio.sleep(1)
+                    continue
+
                 return await self.send(query, timeout=timeout)
             except (FloodWait, FloodPremiumWait) as e:
                 amount = e.value
@@ -454,6 +467,16 @@ class Session:
                     query_name, str(e) or repr(e)
                 )
 
+                if isinstance(e, OSError) and retries > 1:
+                    try:
+                        await self.restart()
+                    except Exception as restart_error:
+                        log.warning(
+                            "[%s] Failed to restart session: %s",
+                            self.client.name,
+                            str(restart_error) or repr(restart_error),
+                        )
+
                 await asyncio.sleep(0.5)
 
         raise TimeoutError("Exceeded maximum number of retries")