Fix WebSocket connection and add error handling

app.py

@@ -100,7 +100,6 @@ def get_stock_history(symbol: str, period: str = "1mo") -> str:
         logger.error(f"Error getting stock history: {str(e)}")
         return f"Error getting stock history: {str(e)}"
 
-# Create tools list
 logger.info("Creating tools list...")
 tools = [get_stock_price, get_stock_info, get_stock_history]
 
@@ -178,9 +177,10 @@ async def websocket_endpoint(websocket: WebSocket):
     chat_history = []  # Initialize chat history
     try:
         while True:
-            data = await websocket.receive_text()
-            logger.info(f"Received message: {data}")
             try:
+                data = await websocket.receive_text()
+                logger.info(f"Received message: {data}")
+                
                 # Add user message to chat history
                 chat_history.append(HumanMessage(content=data))
                 
@@ -198,9 +198,14 @@ async def websocket_endpoint(websocket: WebSocket):
             except Exception as e:
                 logger.error(f"Error processing message: {str(e)}")
                 await websocket.send_json({"type": "error", "content": f"Error processing message: {str(e)}"})
+    except WebSocketDisconnect:
+        logger.info("Client disconnected")
     except Exception as e:
         logger.error(f"WebSocket error: {str(e)}")
-        await websocket.close()
+        try:
+            await websocket.close()
+        except:
+            pass
 
 
 # Serve the HTML frontend

index.html

@@ -104,11 +104,32 @@
     </div>
 
     <script>
-        const ws = new WebSocket('ws://' + window.location.host + '/ws');
+        const ws = new WebSocket('wss://' + window.location.host + '/ws');
         const chatContainer = document.getElementById('chatContainer');
         const userInput = document.getElementById('userInput');
 
+        ws.onopen = function() {
+            console.log('WebSocket connection established');
+        };
+
+        ws.onerror = function(error) {
+            console.error('WebSocket error:', error);
+            const messageDiv = document.createElement('div');
+            messageDiv.className = 'message error-message';
+            messageDiv.textContent = 'Connection error. Please refresh the page.';
+            chatContainer.appendChild(messageDiv);
+        };
+
+        ws.onclose = function() {
+            console.log('WebSocket connection closed');
+            const messageDiv = document.createElement('div');
+            messageDiv.className = 'message error-message';
+            messageDiv.textContent = 'Connection closed. Please refresh the page.';
+            chatContainer.appendChild(messageDiv);
+        };
+
         ws.onmessage = function(event) {
+            console.log('Received message:', event.data);
             const data = JSON.parse(event.data);
             const messageDiv = document.createElement('div');
             messageDiv.className = 'message ' + (data.type === 'ai_message' ? 'ai-message' : 'error-message');
@@ -120,14 +141,21 @@
         function sendMessage() {
             const message = userInput.value.trim();
             if (message) {
-                const messageDiv = document.createElement('div');
-                messageDiv.className = 'message user-message';
-                messageDiv.textContent = message;
-                chatContainer.appendChild(messageDiv);
-                chatContainer.scrollTop = chatContainer.scrollHeight;
+                if (ws.readyState === WebSocket.OPEN) {
+                    const messageDiv = document.createElement('div');
+                    messageDiv.className = 'message user-message';
+                    messageDiv.textContent = message;
+                    chatContainer.appendChild(messageDiv);
+                    chatContainer.scrollTop = chatContainer.scrollHeight;
 
-                ws.send(message);
-                userInput.value = '';
+                    ws.send(message);
+                    userInput.value = '';
+                } else {
+                    const messageDiv = document.createElement('div');
+                    messageDiv.className = 'message error-message';
+                    messageDiv.textContent = 'Not connected. Please refresh the page.';
+                    chatContainer.appendChild(messageDiv);
+                }
             }
         }