Fix real-time messaging

frontend/src/ui/components/chat/MessagePanelRenderer.tsx

@@ -2,6 +2,7 @@ import { useState, useEffect, useRef } from "react";
 import { MessagePanel, type MessagePanelState } from "../../panels/MessagePanel";
 import { ChatMessages } from "./ChatMessages";
 import { ChatInputWrapper } from "./ChatInputWrapper";
+import { setGlobalMessageHandler } from "../../../websocket";
 import defaultAvatar from "../../../resources/images/default-avatar.png";
 
 interface MessagePanelRendererProps {
@@ -26,10 +27,24 @@ export function MessagePanelRenderer({ panel, isChatSwitching }: MessagePanelRen
             };
             
             // Store the handler for cleanup
-            (panel as any).onStateChange = handleStateChange;
+            panel.onStateChange = handleStateChange;
+            
+            // Set up WebSocket message handler for this panel
+            if (panel.handleWebSocketMessage) {
+                setGlobalMessageHandler(panel.handleWebSocketMessage);
+            }
         } else {
             setPanelState(null);
+            // Clear global message handler when no panel is active
+            setGlobalMessageHandler(null);
         }
+        
+        // Cleanup function
+        return () => {
+            if (panel && (panel as any).onStateChange) {
+                (panel as any).onStateChange = null;
+            }
+        };
     }, [panel]);
 
     // Handle chat switching animation

frontend/src/ui/hooks/useChat.ts

@@ -70,40 +70,8 @@ export function useChat() {
         }
     }, [user.authToken]);
 
-    // Handle WebSocket messages
+    // WebSocket messages are now handled by the active panel
-    useEffect(() => {
+    // No need for duplicate handling here
-        const handleWebSocketMessage = (event: MessageEvent) => {
-            try {
-                const response: WebSocketMessage = JSON.parse(event.data);
-                
-                switch (response.type) {
-                    case 'messageEdited':
-                        if (response.data) {
-                            updateMessage(response.data.id, response.data);
-                        }
-                        break;
-                    case 'messageDeleted':
-                        if (response.data && response.data.message_id) {
-                            removeMessage(response.data.message_id);
-                        }
-                        break;
-                    case 'newMessage':
-                        if (response.data) {
-                            addMessage(response.data);
-                        }
-                        break;
-                }
-            } catch (error) {
-                console.error("Error parsing WebSocket message:", error);
-            }
-        };
-
-        websocket.addEventListener("message", handleWebSocketMessage);
-
-        return () => {
-            websocket.removeEventListener("message", handleWebSocketMessage);
-        };
-    }, [addMessage, updateMessage, removeMessage, user.currentUser]);
 
     // Load messages only once when component mounts and user is authenticated
     useEffect(() => {

frontend/src/ui/panels/MessagePanel.ts

@@ -22,7 +22,7 @@ export interface MessagePanelCallbacks {
 export abstract class MessagePanel {
     protected state: MessagePanelState;
     protected callbacks: MessagePanelCallbacks;
-    protected onStateChange: (state: MessagePanelState) => void;
+    public onStateChange: (state: MessagePanelState) => void;
     protected currentUser: UserState;
 
     constructor(
@@ -51,6 +51,9 @@ export abstract class MessagePanel {
     abstract sendMessage(content: string): Promise<void>;
     abstract isDm(): boolean;
     
+    // Optional WebSocket message handler (can be overridden by subclasses)
+    handleWebSocketMessage?: (response: any) => void;
+
     // Common methods
     protected updateState(updates: Partial<MessagePanelState>): void {
         this.state = { ...this.state, ...updates };

frontend/src/ui/state.ts

@@ -1,6 +1,6 @@
 import { create } from "zustand";
 import type { Message, User, WebSocketMessage } from "../core/types";
-import { request } from "../websocket";
+import { request, reconnectWebSocket } from "../websocket";
 import { MessagePanel } from "./panels/MessagePanel";
 import { PublicChatPanel } from "./panels/PublicChatPanel";
 import { DMPanel, type DMPanelData } from "./panels/DMPanel";
@@ -151,6 +151,9 @@ export const useAppState = create<AppState>((set, get) => ({
             }
         }));
 
+        // Reconnect WebSocket with new auth token
+        reconnectWebSocket();
+
         try {
             const payload: WebSocketMessage = {
                 type: "ping",

frontend/src/websocket.ts

@@ -29,6 +29,20 @@ function create(): WebSocket {
  */
 export let websocket: WebSocket = create();
 
+/**
+ * Global WebSocket message handler reference
+ * This will be set by the active panel to handle incoming messages
+ */
+let globalMessageHandler: ((response: WebSocketMessage) => void) | null = null;
+
+/**
+ * Set the global WebSocket message handler
+ * @param handler - Function to handle WebSocket messages
+ */
+export function setGlobalMessageHandler(handler: ((response: WebSocketMessage) => void) | null): void {
+    globalMessageHandler = handler;
+}
+
 export function request(payload: WebSocketMessage): Promise<WebSocketMessage> {
     return new Promise((resolve, reject) => {
         let listener: ((e: MessageEvent) => void) | null = null;
@@ -65,11 +79,28 @@ async function onError() {
     websocket.addEventListener("error", onError);
 }
 
+/**
+ * Recreate WebSocket connection (useful when user logs in)
+ */
+export function reconnectWebSocket(): void {
+    websocket = create();
+    websocket.addEventListener("error", onError);
+}
+
 // --------------
 // Initialization
 // --------------
 
 websocket.addEventListener("message", (e) => {
-    // handleWebSocketMessage(JSON.parse(e.data));
+    try {
+        const response: WebSocketMessage = JSON.parse(e.data);
+        
+        // Route message to global handler if set
+        if (globalMessageHandler) {
+            globalMessageHandler(response);
+        }
+    } catch (error) {
+        console.error("Error parsing WebSocket message:", error);
+    }
 });
 websocket.addEventListener("error", onError);