up

server.py

@@ -487,16 +487,25 @@ async def root_handler(request: web.Request) -> web.Response:
     return web.Response(text=html_content, content_type='text/html')
 
 async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
+    """Handle WebSocket connections with robust error handling"""
     logger.info(f"WebSocket connection attempt - PATH: {request.path}, QUERY: {request.query_string}")
-    logger.info(f"WebSocket request headers: {dict(request.headers)}")
+    
+    # Log request headers at debug level only (could contain sensitive information)
+    logger.debug(f"WebSocket request headers: {dict(request.headers)}")
+    
+    # Prepare a WebSocket response with appropriate settings
+    ws = web.WebSocketResponse(
+        max_msg_size=1024*1024*10,  # 10MB max message size
+        timeout=60.0,
+        heartbeat=30.0  # Add heartbeat to keep connection alive
+    )
+    
+    # Check if WebSocket protocol is supported
+    if not ws.can_prepare(request):
+        logger.error("Cannot prepare WebSocket: WebSocket protocol not supported")
+        return web.Response(status=400, text="WebSocket protocol not supported")
     
     try:
-        ws = web.WebSocketResponse(
-            max_msg_size=1024*1024*10,  # 10MB max message size
-            timeout=60.0,
-            heartbeat=30.0  # Add heartbeat to keep connection alive
-        )
-        
         logger.info("Preparing WebSocket connection...")
         await ws.prepare(request)
         
@@ -510,26 +519,49 @@ async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
         else:
             client_ip = request.headers.get('X-Forwarded-For', 'unknown').split(',')[0].strip()
         
-        logger.info(f"Client {user_id} connecting from IP: {client_ip} - WebSocket ready state: {ws.ready_state}")
+        # Log connection success
+        logger.info(f"Client {user_id} connecting from IP: {client_ip} - WebSocket connection established")
         
-        # Store the user ID in the websocket for easy access
-        ws.user_id = user_id
+        # Mark that the session is established
+        is_session_created = False
         
-        # Create a new session for this user
-        logger.info(f"Creating game session for user {user_id}")
-        user_session = await game_manager.create_session(user_id, ws)
-        logger.info(f"Game session created for user {user_id}")
+        try:
+            # Store the user ID in the websocket for easy access
+            ws.user_id = user_id
+            
+            # Create a new session for this user
+            logger.info(f"Creating game session for user {user_id}")
+            user_session = await game_manager.create_session(user_id, ws)
+            is_session_created = True
+            logger.info(f"Game session created for user {user_id}")
+        except Exception as session_error:
+            logger.error(f"Error creating game session: {str(session_error)}", exc_info=True)
+            if not ws.closed:
+                await ws.close(code=1011, message=f"Server error: {str(session_error)}".encode())
+            if is_session_created:
+                await game_manager.delete_session(user_id)
+            return ws
     except Exception as e:
         logger.error(f"Error establishing WebSocket connection: {str(e)}", exc_info=True)
-        raise
+        if not ws.closed and ws.prepared:
+            await ws.close(code=1011, message=f"Server error: {str(e)}".encode())
+        return ws
     
     # Send initial welcome message
-    await ws.send_json({
-        'action': 'welcome',
-        'userId': user_id,
-        'message': 'Welcome to the MatrixGame WebSocket server!',
-        'scenes': game_manager.valid_scenes
-    })
+    try:
+        await ws.send_json({
+            'action': 'welcome',
+            'userId': user_id,
+            'message': 'Welcome to the MatrixGame WebSocket server!',
+            'scenes': game_manager.valid_scenes
+        })
+        logger.info(f"Sent welcome message to user {user_id}")
+    except Exception as welcome_error:
+        logger.error(f"Error sending welcome message: {str(welcome_error)}")
+        if not ws.closed:
+            await ws.close(code=1011, message=b"Failed to send welcome message")
+        await game_manager.delete_session(user_id)
+        return ws
 
     try:
         async for msg in ws:
@@ -538,6 +570,8 @@ async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
                     data = json.loads(msg.data)
                     action = data.get('action')
                     
+                    logger.debug(f"Received {action} message from user {user_id}")
+                    
                     if action == 'ping':
                         # Respond to ping immediately
                         await ws.send_json({
@@ -551,25 +585,46 @@ async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
                         
                 except json.JSONDecodeError:
                     logger.error(f"Invalid JSON from user {user_id}: {msg.data}")
-                    await ws.send_json({
-                        'error': 'Invalid JSON message',
-                        'success': False
-                    })
+                    if not ws.closed:
+                        await ws.send_json({
+                            'error': 'Invalid JSON message',
+                            'success': False
+                        })
                 except Exception as e:
                     logger.error(f"Error processing WebSocket message for user {user_id}: {str(e)}")
-                    await ws.send_json({
-                        'action': data.get('action') if 'data' in locals() else 'unknown',
-                        'success': False,
-                        'error': f'Error processing message: {str(e)}'
-                    })
+                    if not ws.closed:
+                        await ws.send_json({
+                            'action': data.get('action') if 'data' in locals() else 'unknown',
+                            'success': False,
+                            'error': f'Error processing message: {str(e)}'
+                        })
                     
-            elif msg.type in (WSMsgType.ERROR, WSMsgType.CLOSE):
+            elif msg.type == WSMsgType.ERROR:
+                logger.error(f"WebSocket error for user {user_id}: {ws.exception()}")
+                break
+                
+            elif msg.type == WSMsgType.CLOSE:
+                logger.info(f"WebSocket close received for user {user_id} (code: {msg.data}, message: {msg.extra})")
                 break
                 
+            elif msg.type == WSMsgType.CLOSING:
+                logger.info(f"WebSocket closing for user {user_id}")
+                break
+                
+            elif msg.type == WSMsgType.CLOSED:
+                logger.info(f"WebSocket already closed for user {user_id}")
+                break
+                
+    except Exception as ws_error:
+        logger.error(f"Unexpected WebSocket error for user {user_id}: {str(ws_error)}", exc_info=True)
     finally:
         # Cleanup session
-        await game_manager.delete_session(user_id)
-        logger.info(f"Connection closed for user {user_id}")
+        try:
+            logger.info(f"Cleaning up session for user {user_id}")
+            await game_manager.delete_session(user_id)
+            logger.info(f"Connection closed for user {user_id}")
+        except Exception as cleanup_error:
+            logger.error(f"Error during session cleanup for user {user_id}: {str(cleanup_error)}")
     
     return ws