WebSocket Chat Example

# app.py
import json
from datetime import datetime
from flaxon import Flaxon, WebSocket, WebSocketDisconnect

app = Flaxon("chat-app", debug=True)

# In-memory storage
rooms = {}
users = {}
message_history = {}
max_messages_per_room = 100

@app.websocket("/ws/chat/<room_id>")
async def chat(socket: WebSocket, room_id: str):
    username = socket.scope.get("query_string", "").split("=")[1] if "user=" in socket.scope.get("query_string", "") else "Anonymous"
    user_id = f"{username}-{id(socket)}"

    await socket.accept()
    await socket.join(room_id)

    if room_id not in rooms:
        rooms[room_id] = []
        message_history[room_id] = []

    rooms[room_id].append(user_id)
    users[user_id] = {"username": username, "room": room_id, "socket": socket}

    # Notify others
    await socket.broadcast_json(room_id, {
        "type": "system",
        "event": "user_joined",
        "user": username,
        "users": rooms[room_id],
        "timestamp": datetime.now().isoformat(),
    })

    # Send message history
    await socket.send_json({
        "type": "system",
        "event": "history",
        "messages": message_history[room_id][-max_messages_per_room:],
    })

    try:
        async for message in socket.iter_json():
            if message.get("type") == "message":
                msg = {
                    "type": "message",
                    "user": username,
                    "content": message.get("content"),
                    "timestamp": datetime.now().isoformat(),
                }
                message_history[room_id].append(msg)
                await socket.broadcast_json(room_id, msg)

            elif message.get("type") == "typing":
                await socket.broadcast_json(room_id, {
                    "type": "typing",
                    "user": username,
                    "is_typing": message.get("is_typing", False),
                })

    except WebSocketDisconnect:
        pass
    finally:
        users.pop(user_id, None)
        if room_id in rooms:
            rooms[room_id] = [u for u in rooms[room_id] if u != user_id]

        await socket.leave(room_id)

        await socket.broadcast_json(room_id, {
            "type": "system",
            "event": "user_left",
            "user": username,
            "users": rooms.get(room_id, []),
        })

@app.get("/rooms")
async def list_rooms():
    return {
        "rooms": [
            {
                "name": room_id,
                "users": len(users_in_room),
                "messages": len(message_history.get(room_id, [])),
            }
            for room_id, users_in_room in rooms.items()
        ]
    }

@app.get("/rooms/<room_id>/messages")
async def get_messages(room_id: str):
    return {"messages": message_history.get(room_id, [])}

@app.get("/rooms/<room_id>/users")
async def get_users(room_id: str):
    if room_id in rooms:
        usernames = [users.get(u_id, {}).get("username", "Unknown") for u_id in rooms[room_id]]
        return {"users": usernames}
    return {"users": []}