Update server.py

server.py

@@ -39,23 +39,21 @@ async def handle_sse(request: Request):
                 async for session_message in write_stream_reader:
                     # Handle both dict and SessionMessage objects
                     if isinstance(session_message, dict):
-                        message_data = json.dumps(session_message)
+                        event = session_message.get("event", "message")
+                        data = session_message.get("data", json.dumps(session_message))
                     else:
-                        message_data = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
-                    message = json.loads(message_data)
-                    if not session_id and message.get("event") == "endpoint":
-                        endpoint_url = message.get("data", "")
+                        event = "message"
+                        data = session_message.message.model_dump_json(by_alias=True, exclude_none=True)
+                    message = json.loads(data) if isinstance(data, str) else data
+                    if not session_id and event == "endpoint":
+                        endpoint_url = data
                         if "session_id=" in endpoint_url:
                             session_id = endpoint_url.split("session_id=")[1]
                             placeholder_id = f"placeholder_{id(write_stream)}"
                             if placeholder_id in write_streams:
                                 write_streams[session_id] = write_streams.pop(placeholder_id)
                                 logger.info(f"Updated placeholder {placeholder_id} to session_id {session_id}")
-                    # Send the event as a raw SSE event
-                    await sse_stream_writer.send({
-                        "event": message.get("event", "message"),
-                        "data": message.get("data", message_data)
-                    })
+                    await sse_stream_writer.send({"event": event, "data": data})
 
         sse_stream_writer, sse_stream_reader = anyio.create_memory_object_stream(0)
         async with transport.connect_sse(request.scope, request.receive, request._send) as (read_stream, write_stream):