From 3549f9869100a78b663ff6f30ffeb585987e5f55 Mon Sep 17 00:00:00 2001
From: denis0001-dev <denis0001.dev@ya.ru>
Date: Wed, 26 Nov 2025 16:32:46 +0300
Subject: [PATCH] Ensure ping after auth

---
 frontend/src/core/websocket.ts        | 51 +++++++++++++++++++++++++++
 frontend/src/pages/auth/LoginForm.tsx |  8 +++++
 2 files changed, 59 insertions(+)

diff --git a/frontend/src/core/websocket.ts b/frontend/src/core/websocket.ts
index 0915d19..388c493 100644
--- a/frontend/src/core/websocket.ts
+++ b/frontend/src/core/websocket.ts
@@ -342,4 +342,55 @@ export function request<Request, Response = any>(payload: WebSocketMessage<Reque
 // Initialization
 // --------------
 
+/**
+ * Ensure WebSocket is connected and authenticated after login
+ * This should be called after successful authentication
+ */
+export async function ensureAuthenticated(): Promise<void> {
+    const token = getAuthToken();
+    if (!token) {
+        return;
+    }
+
+    // If WebSocket is not connected, wait for it to connect
+    if (websocket.readyState === WebSocket.CONNECTING) {
+        await new Promise<void>((resolve) => {
+            const checkConnection = () => {
+                if (websocket.readyState === WebSocket.OPEN) {
+                    resolve();
+                } else if (websocket.readyState === WebSocket.CLOSED) {
+                    // Connection failed, try to reconnect
+                    reconnect().then(() => {
+                        setTimeout(checkConnection, 100);
+                    });
+                } else {
+                    setTimeout(checkConnection, 100);
+                }
+            };
+            checkConnection();
+        });
+    } else if (websocket.readyState === WebSocket.CLOSED) {
+        // Reconnect if closed
+        await reconnect();
+    }
+
+    // If WebSocket is open, send ping to authenticate
+    if (websocket.readyState === WebSocket.OPEN) {
+        try {
+            const credentials = {
+                scheme: "Bearer",
+                credentials: token
+            };
+            
+            await request({
+                type: "ping",
+                credentials,
+                data: {}
+            });
+        } catch (error) {
+            console.error("Failed to send ping after login:", error);
+        }
+    }
+}
+
 setupEventHandlers();
\ No newline at end of file
diff --git a/frontend/src/pages/auth/LoginForm.tsx b/frontend/src/pages/auth/LoginForm.tsx
index 8e1410d..17f96c1 100644
--- a/frontend/src/pages/auth/LoginForm.tsx
+++ b/frontend/src/pages/auth/LoginForm.tsx
@@ -12,6 +12,7 @@ import { isElectron } from "@/core/electron/electron";
 import type { Alert, AlertType } from "./Auth";
 import { AuthHeader, AlertsContainer } from "./Auth";
 import styles from "./auth.module.scss";
+import { ensureAuthenticated } from "@/core/websocket";
 
 const loginFieldVariants: Variants = {
     initial: {
@@ -95,6 +96,13 @@ export function LoginForm({ onSwitchMode }: LoginFormProps) {
                     console.error("Key setup failed:", e);
                 }
 
+                // Ensure WebSocket is connected and authenticated
+                try {
+                    await ensureAuthenticated();
+                } catch (e) {
+                    console.error("WebSocket authentication failed:", e);
+                }
+
                 navigate("/chat");
 
                 try {