taleweave-ai/adventure/server.py

import asyncio
from collections import deque
from json import dumps, loads
from logging import getLogger
from threading import Thread
from typing import Dict, Tuple

import websockets

from adventure.context import get_actor_agent_for_name
from adventure.models import Actor, Room, World
from adventure.player import RemotePlayer
from adventure.state import snapshot_world, world_json

logger = getLogger(__name__)

connected = set()
characters: Dict[str, RemotePlayer] = {}
recent_events = deque(maxlen=100)
recent_world = None


async def handler(websocket):
    logger.info("Client connected")
    connected.add(websocket)

    async def next_turn(character: str, prompt: str) -> None:
        await websocket.send(connected, dumps({
            "type": "turn",
            "character": character,
            "prompt": prompt,
        }))

    def sync_turn(character: str, prompt: str) -> bool:
        if websocket not in characters:
            return False

        asyncio.run(next_turn(character, prompt))
        return True

    try:
        if recent_world:
            await websocket.send(recent_world)

        for message in recent_events:
            await websocket.send(message)
    except Exception:
        logger.exception("Failed to send recent messages to new client")

    while True:
        try:
            # if this socket is attached to a character and that character's turn is active, wait for input
            message = await websocket.recv()
            logger.info(f"Received message: {message}")

            try:
                data = loads(message)
                if "become" in data:
                    character = characters.get(websocket)
                    if character:
                        del characters[websocket]

                    character_name = data["become"]
                    actor, _ = get_actor_agent_for_name(character_name)
                    if not actor:
                        logger.error(f"Failed to find actor {character_name}")
                        continue

                    if character_name in [player.name for player in characters.values()]:
                        logger.error(f"Character {character_name} is already in use")
                        continue

                    characters[websocket] = RemotePlayer(actor.name, actor.backstory, sync_turn)
                    logger.info(f"Client {websocket} is now character {character_name}")
                elif websocket in characters:
                    player = characters[websocket]
                    player.input_queue.put(message)

            except Exception:
                logger.exception("Failed to parse message")
        except websockets.ConnectionClosedOK:
            break

    connected.remove(websocket)

    # TODO: swap out the character for the original agent
    if websocket in characters:
        del characters[websocket]

    logger.info("Client disconnected")


socket_thread = None
static_thread = None


def server_json(obj):
    if isinstance(obj, Actor):
        return obj.name

    if isinstance(obj, Room):
        return obj.name

    return world_json(obj)


def send_and_append(message):
    json_message = dumps(message, default=server_json)
    recent_events.append(json_message)
    websockets.broadcast(connected, json_message)
    return json_message


def launch_server():
    global socket_thread, static_thread

    def run_sockets():
        asyncio.run(server_main())

    socket_thread = Thread(target=run_sockets)
    socket_thread.start()


async def server_main():
    async with websockets.serve(handler, "", 8001):
        logger.info("Server started")
        await asyncio.Future()  # run forever


def server_system(world: World, step: int):
    global recent_world
    json_state = {
        **snapshot_world(world, step),
        "type": "world",
    }
    recent_world = send_and_append(json_state)


def server_result(room: Room, actor: Actor, action: str):
    json_action = {
        "actor": actor,
        "result": action,
        "room": room,
        "type": "result",
    }
    send_and_append(json_action)


def server_action(room: Room, actor: Actor, message: str):
    json_input = {
        "actor": actor,
        "input": message,
        "room": room,
        "type": "action",
    }
    send_and_append(json_input)


def server_event(message: str):
    json_broadcast = {
        "message": message,
        "type": "event",
    }
    send_and_append(json_broadcast)