Enhance TelegramProvider and MTProtoService for improved connection management

- Introduced exponential backoff for connection retries in TelegramProvider, enhancing resilience during connection failures.
- Refactored MTProtoService to manage concurrent initialization and connection calls, preventing duplicate requests.
- Updated Redux store configuration to reset volatile fields on rehydrate, ensuring clean state management for Telegram sessions.
- Enhanced error handling in Telegram thunks to ensure proper initialization before connection attempts.

This update significantly improves the reliability and user experience of the Telegram connection process, ensuring smoother interactions and better state management.

.gitignore

@@ -30,3 +30,4 @@ dist-ssr
 *.njsproj
 *.sln
 *.sw?
+references/

src/providers/TelegramProvider.tsx

@@ -4,6 +4,7 @@ import {
   createContext,
   useContext,
   ReactNode,
+  useCallback,
 } from "react";
 import { useAppSelector, useAppDispatch } from "../store/hooks";
 import {
@@ -39,11 +40,14 @@ interface TelegramProviderProps {
   children: ReactNode;
 }
 
+const MAX_RETRIES = 5;
+const BASE_DELAY_MS = 1000;
+
 /**
  * TelegramProvider manages the Telegram MTProto connection
  * - Initializes when app-authenticated (JWT), or has Telegram session / authenticated
  * - Starts init+connect in parallel with login (token) so connect modal is ready sooner
- * - Connects automatically
+ * - Connects automatically with exponential backoff on failure
  * - Provides Telegram context to children
  */
 const TelegramProvider = ({ children }: TelegramProviderProps) => {
@@ -54,22 +58,30 @@ const TelegramProvider = ({ children }: TelegramProviderProps) => {
   const isInitialized = useAppSelector(selectIsInitialized);
   const connectionStatus = useAppSelector(selectConnectionStatus);
   const sessionString = useAppSelector(selectSessionString);
-  const initializedRef = useRef(false);
-  const connectedRef = useRef(false);
+
   const setupInProgressRef = useRef(false);
   const setupCompleteRef = useRef(false);
+  const retryCountRef = useRef(0);
+  const retryTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
 
   const hasSession = !!sessionString;
-
   const shouldSetupTelegram = !!token && !!userId;
 
-  // Initialize Telegram when app-authenticated (token), or have session / Telegram auth
+  const clearRetryTimeout = useCallback(() => {
+    if (retryTimeoutRef.current !== null) {
+      clearTimeout(retryTimeoutRef.current);
+      retryTimeoutRef.current = null;
+    }
+  }, []);
+
+  // Initialize and connect Telegram with exponential backoff on failure
   useEffect(() => {
     if (!shouldSetupTelegram) {
-      initializedRef.current = false;
-      connectedRef.current = false;
+      // Reset all state when conditions no longer met
       setupInProgressRef.current = false;
       setupCompleteRef.current = false;
+      retryCountRef.current = 0;
+      clearRetryTimeout();
       return;
     }
 
@@ -77,6 +89,7 @@ const TelegramProvider = ({ children }: TelegramProviderProps) => {
     if (
       setupCompleteRef.current &&
       isInitialized &&
+      mtprotoService.isReady() &&
       connectionStatus === "connected"
     ) {
       return;
@@ -87,48 +100,69 @@ const TelegramProvider = ({ children }: TelegramProviderProps) => {
     }
 
     const setupTelegram = async () => {
-      // Mark setup as in progress
       setupInProgressRef.current = true;
 
       try {
-        // Initialize if not already initialized
-        if (!isInitialized && !initializedRef.current) {
-          initializedRef.current = true;
+        // Stale-state guard: Redux says isInitialized but the service has no client
+        // (happens after page reload — persist restores isInitialized: true but mtprotoService starts fresh)
+        const needsInit = !mtprotoService.isReady();
+
+        if (needsInit) {
           await dispatch(initializeTelegram(userId)).unwrap();
+          // After init, let the effect re-fire to handle connect
           setupInProgressRef.current = false;
           return;
         }
 
-        if (
-          isInitialized &&
-          connectionStatus !== "connected" &&
-          !connectedRef.current
-        ) {
-          connectedRef.current = true;
+        if (connectionStatus !== "connected") {
           await dispatch(connectTelegram(userId)).unwrap();
           setupInProgressRef.current = false;
-          return; // Exit early, will re-run when connectionStatus updates
+          return;
         }
 
-        // Setup complete - mark as done
+        // Setup complete
         setupInProgressRef.current = false;
         setupCompleteRef.current = true;
+        retryCountRef.current = 0;
       } catch (error) {
         console.error("Failed to setup Telegram:", error);
-        initializedRef.current = false;
-        connectedRef.current = false;
         setupInProgressRef.current = false;
         setupCompleteRef.current = false;
+
+        // Exponential backoff
+        retryCountRef.current += 1;
+        if (retryCountRef.current <= MAX_RETRIES) {
+          const delay = BASE_DELAY_MS * Math.pow(2, retryCountRef.current - 1);
+          console.log(
+            `Telegram setup retry ${retryCountRef.current}/${MAX_RETRIES} in ${delay}ms`,
+          );
+          clearRetryTimeout();
+          retryTimeoutRef.current = setTimeout(() => {
+            retryTimeoutRef.current = null;
+            // Re-trigger by calling setupTelegram again
+            // The effect deps won't have changed, so we call directly
+            setupTelegram();
+          }, delay);
+        } else {
+          console.error(
+            `Telegram setup failed after ${MAX_RETRIES} retries. Giving up.`,
+          );
+        }
       }
     };
 
     setupTelegram();
+
+    return () => {
+      clearRetryTimeout();
+    };
   }, [
     shouldSetupTelegram,
     isInitialized,
     connectionStatus,
     dispatch,
     userId,
+    clearRetryTimeout,
   ]);
 
   // Check connection status once when Telegram reports as connected
@@ -141,8 +175,13 @@ const TelegramProvider = ({ children }: TelegramProviderProps) => {
       return;
     }
 
+    // Only check if the service is actually ready
+    if (!mtprotoService.isReady()) {
+      return;
+    }
+
     const uid = userId;
-    const checkConnection = async () => {
+    const doCheck = async () => {
       try {
         await mtprotoService.checkConnection(uid);
       } catch (error) {
@@ -150,7 +189,7 @@ const TelegramProvider = ({ children }: TelegramProviderProps) => {
       }
     };
 
-    checkConnection();
+    doCheck();
   }, [shouldSetupTelegram, isInitialized, connectionStatus, userId]);
 
   const checkConnection = async (): Promise<boolean> => {

src/services/mtprotoService.ts

@@ -18,6 +18,11 @@ class MTProtoService {
   private readonly apiId: number;
   private readonly apiHash: string;
 
+  // In-flight promise guards — concurrent callers await the same promise
+  private initializePromise: Promise<void> | null = null;
+  private connectPromise: Promise<void> | null = null;
+  private checkConnectionPromise: Promise<boolean> | null = null;
+
   private constructor() {
     // Private constructor to enforce singleton
     // Load API credentials from config once
@@ -40,11 +45,24 @@ class MTProtoService {
   /**
    * Initialize the MTProto client with API credentials.
    * Session is stored in Redux (telegram.byUser[userId].sessionString).
+   * Concurrent calls for the same userId await the same in-flight promise.
    */
   async initialize(userId: string): Promise<void> {
     if (this.isInitialized && this.client && this.userId === userId) {
       return;
     }
+    // If already in-flight for the same user, deduplicate
+    if (this.initializePromise && this.userId === userId) {
+      return this.initializePromise;
+    }
+
+    this.initializePromise = this._doInitialize(userId).finally(() => {
+      this.initializePromise = null;
+    });
+    return this.initializePromise;
+  }
+
+  private async _doInitialize(userId: string): Promise<void> {
     if (this.isInitialized && this.userId !== null && this.userId !== userId) {
       await this.clearSessionAndDisconnect(this.userId);
     }
@@ -76,7 +94,8 @@ class MTProtoService {
   }
 
   /**
-   * Connect to Telegram servers
+   * Connect to Telegram servers.
+   * Concurrent calls await the same in-flight promise.
    */
   async connect(): Promise<void> {
     if (!this.client) {
@@ -86,17 +105,27 @@ class MTProtoService {
     }
 
     if (this.isConnected) {
-      console.log("Already connected to Telegram");
       return;
     }
 
+    if (this.connectPromise) {
+      return this.connectPromise;
+    }
+
+    this.connectPromise = this._doConnect().finally(() => {
+      this.connectPromise = null;
+    });
+    return this.connectPromise;
+  }
+
+  private async _doConnect(): Promise<void> {
     try {
-      await this.client.connect();
+      await this.client!.connect();
       this.isConnected = true;
       console.log("Connected to Telegram successfully");
 
       // Save session string if it changed
-      const newSessionString = this.client.session.save() as string | undefined;
+      const newSessionString = this.client!.session.save() as string | undefined;
       if (newSessionString && newSessionString !== this.sessionString) {
         this.sessionString = newSessionString;
         this.saveSession(newSessionString);
@@ -254,11 +283,23 @@ class MTProtoService {
   }
 
   /**
-   * Check connection status and update user online status
-   * This calls getMe() which also updates the user's online status on Telegram
-   * Automatically initializes and connects if needed
+   * Check connection status and update user online status.
+   * This calls getMe() which also updates the user's online status on Telegram.
+   * Automatically initializes and connects if needed.
+   * Concurrent calls await the same in-flight promise.
    */
   async checkConnection(userId?: string): Promise<boolean> {
+    if (this.checkConnectionPromise) {
+      return this.checkConnectionPromise;
+    }
+
+    this.checkConnectionPromise = this._doCheckConnection(userId).finally(() => {
+      this.checkConnectionPromise = null;
+    });
+    return this.checkConnectionPromise;
+  }
+
+  private async _doCheckConnection(userId?: string): Promise<boolean> {
     try {
       if (!this.isInitialized || !this.client) {
         if (!userId) return false;
@@ -326,8 +367,12 @@ class MTProtoService {
     await this.disconnect();
     this.client = undefined;
     this.isInitialized = false;
+    this.isConnected = false;
     this.sessionString = "";
     this.userId = null;
+    this.initializePromise = null;
+    this.connectPromise = null;
+    this.checkConnectionPromise = null;
   }
 
   /**

src/store/index.ts

@@ -2,6 +2,7 @@ import { configureStore } from "@reduxjs/toolkit";
 import {
   persistStore,
   persistReducer,
+  createTransform,
   FLUSH,
   REHYDRATE,
   PAUSE,
@@ -16,6 +17,8 @@ import userReducer from "./userSlice";
 import telegramReducer from "./telegram";
 import { createLogger } from "redux-logger";
 import { IS_DEV } from "../utils/config";
+import type { TelegramRootState, TelegramState } from "./telegram/types";
+import { initialState as telegramInitialState } from "./telegram/types";
 
 // Persist config for auth only
 const authPersistConfig = {
@@ -24,11 +27,39 @@ const authPersistConfig = {
   whitelist: ["token", "isOnboardedByUser"],
 };
 
+// Strip volatile runtime fields from each per-user Telegram state on rehydrate.
+// These fields reflect in-memory MTProto client state and must start fresh on reload.
+const telegramVolatileTransform = createTransform<
+  TelegramRootState["byUser"],
+  TelegramRootState["byUser"]
+>(
+  // inbound (state -> storage): pass through as-is
+  (inboundState) => inboundState,
+  // outbound (storage -> state): reset volatile fields per user
+  (outboundState) => {
+    const cleaned: Record<string, TelegramState> = {};
+    for (const [userId, userState] of Object.entries(outboundState)) {
+      cleaned[userId] = {
+        ...userState,
+        isInitialized: telegramInitialState.isInitialized,
+        connectionStatus: telegramInitialState.connectionStatus,
+        connectionError: telegramInitialState.connectionError,
+        isLoadingChats: telegramInitialState.isLoadingChats,
+        isLoadingMessages: telegramInitialState.isLoadingMessages,
+        isLoadingThreads: telegramInitialState.isLoadingThreads,
+      };
+    }
+    return cleaned;
+  },
+  { whitelist: ["byUser"] },
+);
+
 // Persist config for telegram state (scoped by user in byUser)
 const telegramPersistConfig = {
   key: "telegram",
   storage,
   whitelist: ["byUser"],
+  transforms: [telegramVolatileTransform],
 };
 
 const persistedAuthReducer = persistReducer(authPersistConfig, authReducer);

src/store/telegram/thunks.ts

@@ -28,6 +28,11 @@ export const initializeTelegram = createAsyncThunk(
 export const connectTelegram = createAsyncThunk(
   "telegram/connect",
   async (_userId: string, { rejectWithValue }) => {
+    if (!mtprotoService.isReady()) {
+      return rejectWithValue(
+        "MTProto client not initialized. Call initialize() first.",
+      );
+    }
     try {
       await mtprotoService.connect();
       return true;
@@ -44,6 +49,12 @@ export const connectTelegram = createAsyncThunk(
 export const checkAuthStatus = createAsyncThunk(
   "telegram/checkAuthStatus",
   async (userId: string, { rejectWithValue, getState }) => {
+    if (!mtprotoService.isReady()) {
+      return rejectWithValue(
+        "MTProto client not initialized. Call initialize() first.",
+      );
+    }
+
     const now = Date.now();
     if (isCheckingAuth && now - lastCheckTime < MIN_CHECK_INTERVAL) {
       const state = getState() as RootState;
@@ -92,6 +103,11 @@ export const checkAuthStatus = createAsyncThunk(
 export const fetchChats = createAsyncThunk(
   "telegram/fetchChats",
   async (_userId: string, { rejectWithValue }) => {
+    if (!mtprotoService.isReady()) {
+      return rejectWithValue(
+        "MTProto client not initialized. Call initialize() first.",
+      );
+    }
     try {
       const client = mtprotoService.getClient();
       const dialogs = await mtprotoService.withFloodWaitHandling(async () => {
@@ -117,6 +133,11 @@ export const fetchMessages = createAsyncThunk(
     }: { userId: string; chatId: string; limit?: number; offsetId?: number },
     { rejectWithValue },
   ) => {
+    if (!mtprotoService.isReady()) {
+      return rejectWithValue(
+        "MTProto client not initialized. Call initialize() first.",
+      );
+    }
     try {
       void _userId;
       const client = mtprotoService.getClient();