1 files changed, 29 insertions, 13 deletions
diff --git a/kennel/main.py b/kennel/main.py
index 6f661cb..fc0fd0e 100644
--- a/kennel/main.py
+++ b/kennel/main.py
@@ -11,7 +11,14 @@ from fastapi import (
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 from kennel.engine.systems.system import SystemType
-from kennel.engine.systems.network import Event, Publishable, EventType
+from kennel.engine.systems.network import (
+    Event,
+    Publishable,
+    EventType,
+    InitialStateEvent,
+    SetControllableEvent,
+    EntityBornEvent,
+)
 from typing import Annotated, Optional
 from kennel.kennel import (
     kennel,
@@ -76,7 +83,7 @@ class WebSocketClient(Publishable):
         self.websocket = websocket
 
     async def publish(self, event: Event):
-        await self.websocket.send_json(event.dict())
+        await self.websocket.send_json([event.event_type, event.data])
 
 
 @app.websocket("/ws")
@@ -85,19 +92,16 @@ async def websocket_endpoint(
     session: Annotated[str, Depends(get_cookie_or_token)],
 ):
     await websocket.accept()
+    client = WebSocketClient(websocket)
     logger.info(f"Websocket connection established for session {session}")
 
-    await websocket.send_json(
-        {
-            "event_type": EventType.INITIAL_STATE,
-            "data": {
-                "world": {"width": config.WORLD_WIDTH, "height": config.WORLD_HEIGHT},
-                "entities": kennel.entity_manager.to_dict(),
-            },
-        }
+    initial_state = InitialStateEvent(
+        config.WORLD_WIDTH, config.WORLD_HEIGHT, kennel.entity_manager.to_dict()
     )
+    await client.publish(initial_state)
 
     session_entities = create_session_controllable_entities(session)
+    logger.info(f"Creating {len(session_entities)} entities for session {session}")
     try:
         network_system = system_manager.get_system(SystemType.NETWORK)
         if network_system is None:
@@ -106,14 +110,26 @@ async def websocket_endpoint(
         for entity in session_entities:
             logger.info(f"Adding entity {entity.id} for session {session}")
             entity_manager.add_entity(entity)
-            network_system.add_event(Event(EventType.ENTITY_BORN, {"id": entity.id}))
+            network_system.add_event(EntityBornEvent(entity))
+
+            set_controllable_event = SetControllableEvent(entity.id, session)
+            await client.publish(set_controllable_event)
 
         network_system.add_client(session, WebSocketClient(websocket))
         while True:
             message = await websocket.receive_json()
-            network_system.client_event(session, Event.from_dict(message))
+            if type(message) is not list:
+                message = [message]
+            events = [Event.from_dict(event) for event in message]
+            if not all([event is not None for event in events]):
+                logger.info(f"Invalid events in: {message}"[:100])
+                continue
+            for event in events:
+                network_system.client_event(session, event)
+    except WebSocketDisconnect as e:
+        logger.info(f"Websocket connection closed by client: {session}")
     except Exception as e:
-        logger.error(f"WebSocket exception {e}")
+        logger.error("Exception occurred", exc_info=e)
     finally:
         logger.info("Websocket connection closed")
         for entity in session_entities: