.

main.py

@@ -537,8 +537,18 @@ async def websocket_endpoint(websocket: WebSocket):
                 
                 # Send the generated frame back to the client
                 print(f"[{time.perf_counter():.3f}] Sending image to client...")
-                await websocket.send_json({"image": img_str})
-                print(f"[{time.perf_counter():.3f}] Image sent. Queue size before next_input: {input_queue.qsize()}")
+                try:
+                    await websocket.send_json({"image": img_str})
+                    print(f"[{time.perf_counter():.3f}] Image sent. Queue size before next_input: {input_queue.qsize()}")
+                except RuntimeError as e:
+                    if "Cannot call 'send' once a close message has been sent" in str(e):
+                        print(f"[{time.perf_counter():.3f}] WebSocket closed, skipping image send")
+                        break
+                    else:
+                        raise e
+                except Exception as e:
+                    print(f"[{time.perf_counter():.3f}] Error sending image: {e}")
+                    break
 
                 # Log the input
                 log_interaction(client_id, data, generated_frame=sample_img)
@@ -559,6 +569,12 @@ async def websocket_endpoint(websocket: WebSocket):
                 is_processing = False
                 return
             
+            # Check if WebSocket is still open by checking if it's in a closed state
+            if websocket.client_state.value >= 2:  # CLOSING or CLOSED
+                print(f"[{current_time:.3f}] WebSocket in closed state ({websocket.client_state.value}), stopping processing")
+                is_processing = False
+                return
+            
             #if is_processing:
             #    print(f"[{current_time:.3f}] Already processing an input. Will check again later.")
             #    return
@@ -657,6 +673,11 @@ async def websocket_endpoint(websocket: WebSocket):
                 await input_queue.put(data)
                 print(f"[{receive_time:.3f}] Received input. Queue size now: {input_queue.qsize()}")
                 
+                # Check if WebSocket is still open before processing
+                if websocket.client_state.value >= 2:  # CLOSING or CLOSED
+                    print(f"[{receive_time:.3f}] WebSocket closed, skipping processing")
+                    continue
+                
                 # If we're not currently processing, start processing this input
                 if not is_processing:
                     print(f"[{receive_time:.3f}] Not currently processing, will call process_next_input()")

static/index.html

@@ -93,7 +93,7 @@
         <div class="canvas-container">
             <div id="timeoutWarning" class="alert alert-warning" style="display: none; margin-bottom: 15px;">
                 <strong>Connection Timeout Warning:</strong> 
-                No user activity detected. Connection will be dropped in <span id="timeoutCountdown">10</span> seconds.
+                No user activity detected. Connection will be dropped in <span id="timeoutCountdown">10</span> seconds and page will refresh automatically.
                 <button type="button" class="btn btn-sm btn-primary ms-2" onclick="resetTimeout()">Stay Connected</button>
             </div>
             
@@ -190,11 +190,24 @@
 
             socket.onclose = function(event) {
                 console.log("WebSocket connection closed. Attempting to reconnect...");
+                console.log("Close event code:", event.code, "reason:", event.reason);
+                
+                // Check if this was a timeout closure
+                if (event.code === 1000 && event.reason === "User inactivity timeout") {
+                    console.log("Connection closed due to timeout - refreshing page");
+                    stopAutoInput(); // Stop auto-input when connection is lost
+                    stopTimeoutCountdown(); // Stop timeout countdown when connection is lost
+                    //clearInterval(heartbeatInterval);
+                    // Refresh the page to reconnect
+                    window.location.reload();
+                    return;
+                }
+                
                 isConnected = false;
                 stopAutoInput(); // Stop auto-input when connection is lost
                 stopTimeoutCountdown(); // Stop timeout countdown when connection is lost
-                clearInterval(heartbeatInterval);
-                scheduleReconnection();
+                //clearInterval(heartbeatInterval);
+                //scheduleReconnection();
             };
 
             socket.onerror = function(error) {
@@ -387,8 +400,12 @@
                 }
                 
                 if (timeoutCountdown <= 0) {
+                    stopAutoInput();
                     stopTimeoutCountdown();
-                    console.log("Connection timeout countdown finished");
+                    //clearInterval(heartbeatInterval);
+                    console.log("Connection timeout countdown finished - refreshing page");
+                    // Refresh the page to reconnect
+                    window.location.reload();
                 }
             }, 1000);
         }
@@ -416,12 +433,20 @@
         }
 
         function resetTimeout() {
-            // Send a heartbeat to reset the server's timeout
-            if (socket && socket.readyState === WebSocket.OPEN) {
+            // Send a regular input to reset the server's timeout
+            if (socket && socket.readyState === WebSocket.OPEN && lastSentPosition) {
                 try {
-                    socket.send(JSON.stringify({ type: "heartbeat" }));
+                    socket.send(JSON.stringify({
+                        "x": lastSentPosition.x,
+                        "y": lastSentPosition.y,
+                        "is_left_click": false,
+                        "is_right_click": false,
+                        "keys_down": [],
+                        "keys_up": []
+                    }));
+                    updateLastUserInputTime(); // Update for auto-input mechanism
                 } catch (error) {
-                    console.error("Error sending heartbeat:", error);
+                    console.error("Error sending timeout reset input:", error);
                 }
             }
             stopTimeoutCountdown();