app.py

import asyncio
import json
import os
import secrets
import signal
from collections import defaultdict
from websockets import serve, ConnectionClosed

# Dictionary to store editor content for each user
editor_content = defaultdict(str)

# Set of connected users
connected_users = set()


async def broadcast_changes(websocket, user, change):
    """
    Broadcast code editing changes to all connected clients.
    """
    event = {
        "type": "edit",
        "user": user,
        "change": change,
    }

    # Broadcast the change to all connected clients
    for user_socket in connected_users:
        if user_socket != websocket:
            await user_socket.send(json.dumps(event))


async def handle_code_editing(websocket, user):
    """
    Handle code editing events from a user.
    """
    connected_users.add(websocket)

    try:
        # Send the initial state of the document to the new user
        await websocket.send(json.dumps({"type": "init", "content": editor_content[user]}))

        async for message in websocket:
            event = json.loads(message)
            assert event["type"] == "edit"
            change = event["change"]

            # Apply the change to the document
            editor_content[user] = change["text"]

            # Broadcast the change to all connected clients
            await broadcast_changes(websocket, user, change)

    except ConnectionClosed:
        pass
    finally:
        # Remove the user when the WebSocket connection is closed
        connected_users.remove(websocket)


async def editor_handler(websocket, path):
    """
    Handle a connection and dispatch it according to who is connecting.
    """
    # Receive and parse the "init" event from the UI.
    message = await websocket.recv()
    event = json.loads(message)
    assert event["type"] == "init"

    user = event.get("user", None)

    if not user:
        # Assign a random username for simplicity
        user = secrets.token_urlsafe(8)

    try:
        await handle_code_editing(websocket, user)
    except Exception as e:
        print(f"Error: {e}")
    finally:
        connected_users.remove(websocket)


async def main():
    # Set the stop condition when receiving SIGTERM
    loop = asyncio.get_running_loop()
    stop = loop.create_future()
    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)

    
    async with serve(editor_handler, "0.0.0.0", 7860):
        await stop


if __name__ == "__main__":
    asyncio.run(main())