From be28a11db643ab21ef1994b416f1137dad1ad03f Mon Sep 17 00:00:00 2001
From: vvzvlad <git@vvzvlad.xyz>
Date: Sun, 27 Apr 2025 18:31:29 +0400
Subject: [PATCH] implement connection handling in TelegramClient: add
 disconnect tracking and application restart on multiple disconnections

---
 telegram_client.py | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/telegram_client.py b/telegram_client.py
index 7dd015b..be02181 100644
--- a/telegram_client.py
+++ b/telegram_client.py
@@ -11,10 +11,13 @@
 import logging
 import os
 import asyncio
+import sys
+import signal
 import uvloop
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 from pyrogram import Client
+from pyrogram.handlers import DisconnectHandler
 from config import get_settings
 
 logger = logging.getLogger(__name__)
@@ -29,6 +32,9 @@ class TelegramClient:
             api_hash=settings["tg_api_hash"],
             workdir=settings["session_path"],
         )
+        self.disconnect_count = 0
+        self.max_disconnects = 3
+        self._setup_connection_handlers()
 
     def _ensure_session_directory(self):
         try:
@@ -38,15 +44,48 @@ class TelegramClient:
             logger.error(f'Failed to create session directory {settings["session_path"]}: {str(e)}')
             raise
 
+    def _setup_connection_handlers(self):
+        """Sets up connection/disconnection handlers"""
+        self.client.add_handler(DisconnectHandler(self._on_disconnect))
+        logger.info("connection_handlers: connection handlers set up")
+
+    def _on_disconnect(self, _client):
+        """Handles disconnection from Telegram servers"""
+        self.disconnect_count += 1
+        logger.warning(f"connection_handler: connection lost (#{self.disconnect_count})")
+        
+        if self.disconnect_count >= self.max_disconnects:
+            logger.critical(f"connection_handler: reached disconnect limit ({self.max_disconnects})")
+            self._restart_app()
+
     async def start(self):
         try:
             if not self.client.is_connected:
                 await self.client.start()
                 logger.info("Telegram client connected successfully")
+                # Reset disconnect counter on successful connection
+                self.disconnect_count = 0
+                logger.info("connection_handler: connection established")
         except Exception as e:
             logger.error(f"Failed to start Telegram client: {str(e)}")
             raise
 
+    def _restart_app(self):
+        """Restarts the application"""
+        logger.warning("connection_handler: restarting application")
+        try:
+            # Properly shutdown client before restart
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                loop.create_task(self.stop())
+            
+            # Use os.execv to restart the process with the same arguments
+            os.execv(sys.executable, [sys.executable] + sys.argv)
+        except Exception as e:
+            logger.error(f"connection_handler: error during restart: {str(e)}")
+            # Emergency termination in case of restart failure
+            os.kill(os.getpid(), signal.SIGTERM)
+
     async def stop(self):
         if self.client.is_connected:
             await self.client.stop()