add disconnect mechanism

app.py

@@ -14,6 +14,9 @@ import websockets
 app = FastAPI()
 templates = Jinja2Templates(directory=".")
 
+# Track active WebSocket connections
+active_websockets = set()
+
 # OAuth configuration from HF Spaces environment
 OAUTH_CLIENT_ID = os.getenv("OAUTH_CLIENT_ID")
 OAUTH_CLIENT_SECRET = os.getenv("OAUTH_CLIENT_SECRET")
@@ -420,6 +423,10 @@ async def websocket_video_gen(websocket: WebSocket, user_fal_key: Optional[str]
 
     await websocket.accept()
 
+    # Track this connection
+    active_websockets.add(websocket)
+    print(f"WebSocket connected. Active connections: {len(active_websockets)}")
+
     # Get user from cookie
     access_token = websocket.cookies.get("access_token")
     if not access_token:
@@ -470,7 +477,7 @@ async def websocket_video_gen(websocket: WebSocket, user_fal_key: Optional[str]
     
     # Connect to FAL WebSocket
     fal_ws_url = f"wss://fal.run/fal-ai/krea-wan-14b/ws?fal_jwt_token={fal_token}"
-    
+
     try:
         async with websockets.connect(fal_ws_url) as fal_ws:
             # Relay messages between client and FAL
@@ -483,7 +490,7 @@ async def websocket_video_gen(websocket: WebSocket, user_fal_key: Optional[str]
                         await fal_ws.send(data)
                 except Exception as e:
                     print(f"Client to FAL error: {e}")
-            
+
             async def fal_to_client():
                 try:
                     while True:
@@ -496,7 +503,7 @@ async def websocket_video_gen(websocket: WebSocket, user_fal_key: Optional[str]
                             await websocket.send_bytes(message)
                 except Exception as e:
                     print(f"FAL to client error: {e}")
-            
+
             # Run both directions concurrently
             import asyncio
             await asyncio.gather(
@@ -505,4 +512,8 @@ async def websocket_video_gen(websocket: WebSocket, user_fal_key: Optional[str]
             )
     except Exception as e:
         print(f"WebSocket proxy error: {e}")
-        await websocket.close(code=1011, reason=str(e))
+        await websocket.close(code=1011, reason=str(e))
+    finally:
+        # Remove from active connections
+        active_websockets.discard(websocket)
+        print(f"WebSocket disconnected. Active connections: {len(active_websockets)}")

index.html

@@ -1210,6 +1210,8 @@
             promptUpdateTimer: null,
             pendingPromptUpdate: null,
             generationFinished: false,
+            idleTimeout: null,
+            lastFrameTime: null,
 
             init() {
                 this.setupEventListeners();
@@ -1449,6 +1451,7 @@
 
                         this.ws.onopen = () => {
                             this.showInfo('Connected! Waiting for ready signal...');
+                            this.startIdleTimeout();
                         };
 
                         this.ws.onmessage = async (event) => {
@@ -1590,6 +1593,9 @@
                 this.frameCount++;
                 document.getElementById('frameCount').textContent = this.frameCount;
 
+                // Reset idle timeout on every frame
+                this.startIdleTimeout();
+
                 if (this.frameCount === 1) {
                     document.getElementById('spinner').classList.add('hidden');
                     this.startPlaybackLoop();
@@ -1893,7 +1899,31 @@
                 }
             },
 
+            startIdleTimeout() {
+                // Clear existing timeout
+                if (this.idleTimeout) {
+                    clearTimeout(this.idleTimeout);
+                }
+
+                // Set new timeout - disconnect after 15 seconds of no frames
+                this.idleTimeout = setTimeout(() => {
+                    if (this.isGenerating) {
+                        this.showError('The FAL API server is too busy, try again soon!');
+                        this.disconnect();
+                    }
+                }, 15000);
+            },
+
+            clearIdleTimeout() {
+                if (this.idleTimeout) {
+                    clearTimeout(this.idleTimeout);
+                    this.idleTimeout = null;
+                }
+            },
+
             disconnect() {
+                this.clearIdleTimeout();
+
                 if (this.ws) {
                     this.ws.close();
                     this.ws = null;