Update server.py

server.py

@@ -11,13 +11,13 @@ app = FastAPI()
 # Configure CORS
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Allows all origins
+    allow_origins=["*"],
     allow_credentials=True,
-    allow_methods=["*"],  # Allows all methods
-    allow_headers=["*"],  # Allows all headers
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
-# Serve static files from the "static" directory
+# Serve static files
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 # Set up Jinja2 templates
@@ -74,7 +74,8 @@ class ConnectionManager:
             if len(self.groups[group_name]) == 0:
                 del self.groups[group_name]
                 print(f"Group '{group_name}' deleted (no members)")
-        del self.connections[websocket]
+        if websocket in self.connections:
+            del self.connections[websocket]
         print(f"Client removed from group '{group_name}'")
 
     async def broadcast_to_group(self, message: str, group_name: str):
@@ -94,16 +95,20 @@ async def websocket_endpoint(websocket: WebSocket):
             data = await websocket.receive_text()
             try:
                 message = json.loads(data)
-                if message['type'] == 'join':
-                    group_name = message['group']
-                    await manager.join_group(websocket, group_name)
-                    # Notify others in group about new member
-                    notification = json.dumps({
-                        "type": "system",
-                        "message": "New member joined"
-                    })
-                    await manager.broadcast_to_group(notification, group_name)
-                elif message['type'] == 'message':
+                # Handle join messages
+                if message.get('type') == 'join':
+                    group_name = message.get('group')
+                    if group_name:
+                        await manager.join_group(websocket, group_name)
+                    else:
+                        error = json.dumps({
+                            "type": "error",
+                            "message": "Missing group name in join request"
+                        })
+                        await websocket.send_text(error)
+                
+                # Broadcast all other messages to the group
+                else:
                     if websocket in manager.connections:
                         group_name = manager.connections[websocket]
                         await manager.broadcast_to_group(data, group_name)
@@ -113,29 +118,15 @@ async def websocket_endpoint(websocket: WebSocket):
                             "message": "Join a group before sending messages"
                         })
                         await websocket.send_text(error)
+            
             except json.JSONDecodeError:
-                error = json.dumps({
-                    "type": "error",
-                    "message": "Invalid JSON format"
-                })
-                await websocket.send_text(error)
-            except KeyError:
-                error = json.dumps({
-                    "type": "error",
-                    "message": "Missing required fields in message"
-                })
-                await websocket.send_text(error)
+                # Handle non-JSON messages by broadcasting directly
+                if websocket in manager.connections:
+                    group_name = manager.connections[websocket]
+                    await manager.broadcast_to_group(data, group_name)
+    
     except WebSocketDisconnect:
-        # Get group before disconnecting
-        group_name = manager.connections.get(websocket)
         manager.disconnect(websocket)
-        # Notify group about disconnection
-        if group_name:
-            notification = json.dumps({
-                "type": "system",
-                "message": "A member has left"
-            })
-            await manager.broadcast_to_group(notification, group_name)
 
 
 if __name__ == "__main__":