Files
VDM/server/main.py
T

423 lines
18 KiB
Python

# server/main.py
import base64
import asyncio
from pathlib import Path
from typing import Dict, List, Set
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from starlette.websockets import WebSocketState
from .config import settings
from .database_manager import DatabaseManager
from .models import (
Room,
WSIncomingMessage,
WSOutgoingMessage,
RegisterRequest,
Player,
LoginRequest,
)
from .persistence_manager import PersistenceManager
from .room_manager import RoomManager
from .story_manager import StoryManager
from .audio_manager import AudioManager
from .game_manager import DiceRoller
from .logger import logger
from .user_manager import UserManager
# ===================================================================
# Application Setup
# ===================================================================
BASE_DIR = Path(__file__).resolve().parent.parent
app = FastAPI(title="VDM - Virtual Dungeon Master")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
app.mount("/static", StaticFiles(directory=BASE_DIR / "web"), name="static")
app.mount("/audio", StaticFiles(directory=Path(settings.paths.audio_out_dir)), name="audio")
class ConnectionManager:
"""Manages active WebSocket connections for each room."""
def __init__(self):
self.connections: Dict[str, Set[WebSocket]] = {}
async def connect(self, room_id: str, websocket: WebSocket):
await websocket.accept()
self.connections.setdefault(room_id, set()).add(websocket)
logger.info(f"New connection in room '{room_id}'. Total: {len(self.connections[room_id])}")
def disconnect(self, room_id: str, websocket: WebSocket):
if room_id in self.connections:
self.connections[room_id].discard(websocket)
logger.info(f"Disconnected from room '{room_id}'. Remaining: {len(self.connections.get(room_id, set()))}")
async def broadcast(self, room_id: str, message: WSOutgoingMessage):
if room_id not in self.connections: return
payload = message.model_dump_json()
tasks = [
connection.send_text(payload)
for connection in self.connections.get(room_id, set())
if connection.client_state == WebSocketState.CONNECTED
]
await asyncio.gather(*tasks)
# --- Instantiate Managers ---
db_manager = DatabaseManager(
sessions_db_path=Path(settings.memory.sessions_db_file),
users_db_path=Path(settings.memory.users_db_file)
)
persistence_manager = PersistenceManager(db_manager=db_manager)
user_manager = UserManager(db_manager=db_manager)
room_manager = RoomManager(user_manager=user_manager, persistence_manager=persistence_manager)
story_manager = StoryManager()
audio_manager = AudioManager()
game_manager = DiceRoller()
connection_manager = ConnectionManager()
# ===================================================================
# Core Game Loop Logic
# ===================================================================
async def _start_game_setup_turn(room_id: str):
"""
Handles the very first turn of the game (the GM's setup prompt),
respecting the streaming setting.
"""
room_manager.get_or_create_room(room_id)
if settings.audio.enable_streaming:
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="stream_start", payload={})
)
full_gm_response = ""
text_generator = story_manager.generate_gm_response_stream(room_id, [])
async for text_chunk in text_generator:
if not text_chunk:
continue
full_gm_response += text_chunk
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="chat_chunk", payload={"content": text_chunk})
)
audio_generator = audio_manager.synthesize_stream(text_chunk, room_id)
async for audio_chunk in audio_generator:
encoded_chunk = base64.b64encode(audio_chunk).decode("utf-8")
await connection_manager.broadcast(
room_id,
WSOutgoingMessage(kind="audio_chunk", payload={"chunk": encoded_chunk}),
)
gm_message = room_manager.add_message(
room_id, "gm", "GM", full_gm_response.strip(), audio_url=None
)
await connection_manager.broadcast(
room_id,
WSOutgoingMessage(kind="stream_end", payload={"final_message": gm_message.model_dump()}),
)
else:
gm_prompt = await story_manager.generate_gm_response(room_id, [])
audio_url = await audio_manager.synthesize(gm_prompt, room_id)
gm_message = room_manager.add_message(
room_id, "gm", "GM", gm_prompt, audio_url=audio_url
)
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="chat", payload=gm_message.model_dump())
)
if audio_url:
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="audio", payload={"url": audio_url})
)
async def _advance_turn(room_id: str, submitter: Player):
"""
Orchestrates the GM's turn, handling both streaming and non-streaming modes.
"""
room_state = room_manager.get_room(room_id)
if (
not room_state
or room_state.turn_state == "GM_PROCESSING"
or not room_state.current_turn_actions
):
return
room_state.turn_state = "GM_PROCESSING"
await connection_manager.broadcast(
room_id,
WSOutgoingMessage(
kind="system",
payload={"message": f"{submitter.name} submitted the turn. The GM ponders..."},
),
)
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="state_update", payload=room_state.model_dump())
)
turn_actions = {
room_state.players[pid].name: action
for pid, action in room_state.current_turn_actions.items()
if pid in room_state.players
}
history = [msg.model_dump() for msg in room_state.messages]
consolidated_actions_text = "\n".join(
f"[{name}] {action}" for name, action in turn_actions.items()
)
room_manager.add_message(room_id, "party", "Party Actions", consolidated_actions_text)
if settings.audio.enable_streaming:
await _advance_turn_streaming(room_id, history, turn_actions)
else:
await _advance_turn_non_streaming(room_id, history, turn_actions)
room_state.current_turn_actions.clear()
room_state.turn_state = "WAITING_FOR_ACTIONS"
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="state_update", payload=room_state.model_dump())
)
async def _advance_turn_streaming(
room_id: str, history: List[Dict], turn_actions: Dict[str, str]
):
"""Handles the game turn with real-time streaming of text and audio."""
logger.info(f"Advancing turn for room '{room_id}' with STREAMING.")
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="stream_start", payload={})
)
full_gm_response = ""
text_generator = story_manager.generate_gm_response_stream(
room_id, history, turn_actions
)
async for text_chunk in text_generator:
if not text_chunk:
continue
full_gm_response += text_chunk
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="chat_chunk", payload={"content": text_chunk})
)
audio_generator = audio_manager.synthesize_stream(text_chunk, room_id)
async for audio_chunk in audio_generator:
encoded_chunk = base64.b64encode(audio_chunk).decode("utf-8")
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="audio_chunk", payload={"chunk": encoded_chunk})
)
gm_message = room_manager.add_message(
room_id, "gm", "GM", full_gm_response.strip(), audio_url=None
)
await connection_manager.broadcast(
room_id,
WSOutgoingMessage(kind="stream_end", payload={"final_message": gm_message.model_dump()}),
)
async def _advance_turn_non_streaming(
room_id: str, history: List[Dict], turn_actions: Dict[str, str]
):
"""Handles the game turn by generating the full response before sending."""
logger.info(f"Advancing turn for room '{room_id}' NON-STREAMING.")
gm_response = await story_manager.generate_gm_response(
room_id, history, turn_actions
)
audio_url = await audio_manager.synthesize(gm_response, room_id)
gm_message = room_manager.add_message(
room_id, "gm", "GM", gm_response, audio_url=audio_url
)
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="chat", payload=gm_message.model_dump())
)
if audio_url:
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="audio", payload={"url": audio_url})
)
async def _resume_game_turn(room_id: str, player: Player):
"""Resumes a game non-streamed for simplicity."""
room_state = room_manager.get_room(room_id)
# The check for host actions will now happen inside the websocket_endpoint
if (not room_state or room_state.game_state != "PLAYING"): return
await connection_manager.broadcast(
room_id,
WSOutgoingMessage(kind="system", payload={"message": f"{player.name} is resuming the game..."}),
)
room_state.turn_state = "GM_PROCESSING"
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="state_update", payload=room_state.model_dump())
)
history = [msg.model_dump() for msg in room_state.messages]
gm_summary = await story_manager.generate_resume_summary(room_id, history)
audio_url = await audio_manager.synthesize(gm_summary, room_id)
gm_message = room_manager.add_message(
room_id, "gm", "GM", gm_summary, audio_url=audio_url
)
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="chat", payload=gm_message.model_dump())
)
if audio_url:
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="audio", payload={"url": audio_url})
)
room_state.turn_state = "WAITING_FOR_ACTIONS"
await connection_manager.broadcast(
room_id, WSOutgoingMessage(kind="state_update", payload=room_state.model_dump())
)
# ===================================================================
# API & WebSocket Endpoints
# ===================================================================
@app.get("/")
async def get_root():
return FileResponse(BASE_DIR / "web/index.html")
@app.post("/api/register")
async def register_player(request: RegisterRequest):
success, message = user_manager.register_player(request.name, request.avatar_style, request.password)
if not success: raise HTTPException(status_code=400, detail=message)
return JSONResponse(content={"message": message})
@app.post("/api/login")
async def login_player(request: LoginRequest):
user_data = user_manager.login(request.name, request.password)
if not user_data: raise HTTPException(status_code=401, detail="Invalid username or password.")
return JSONResponse(content=user_data)
@app.get("/api/voices")
async def get_voices():
return JSONResponse(content=audio_manager.list_voices())
@app.get("/favicon.ico", include_in_schema=False)
async def favicon():
return FileResponse(BASE_DIR / "web/favicon.ico")
@app.websocket("/ws/{room_id}/{player_id}/{player_token}")
async def websocket_endpoint(
websocket: WebSocket, room_id: str, player_id: str, player_token: str
):
add_player_result = room_manager.add_player(room_id, player_id, player_token)
if not add_player_result:
await websocket.close(code=4001, reason="Invalid session token.")
return
await connection_manager.connect(room_id, websocket)
room, player = add_player_result
if room.messages:
await websocket.send_text(WSOutgoingMessage(kind="chat_history", payload={"messages": [m.model_dump() for m in room.messages]}).model_dump_json())
# --- Permanent Owner and Host Logic ---
# If the room has no permanent owner, this player becomes the owner.
if not room.owner_username:
room.owner_username = player.name
logger.info(f"Player '{player.name}' is the permanent owner of room '{room_id}'.")
# The active host is always the owner, if they are present.
if room.owner_username == player.name:
room.host_player_id = player.id
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="system", payload={"message": f"{player.name} has joined the game!"}))
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="state_update", payload=room.model_dump()))
try:
while True:
data = await websocket.receive_text()
try:
msg = WSIncomingMessage.model_validate_json(data)
room_state = room_manager.get_room(room_id)
if not room_state: continue
# Check if the current player is the room owner for host actions.
is_owner = (room_state.owner_username == player.name)
if msg.kind == "start_game":
if is_owner and room_state.game_state == "LOBBY":
room_state.game_state = "PLAYING"
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="system", payload={"message": "The game is starting..."}))
await _start_game_setup_turn(room_id)
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="state_update", payload=room_state.model_dump()))
elif msg.kind == "resume_game":
if is_owner:
await _resume_game_turn(room_id, player)
elif msg.kind == "submit_turn":
await _advance_turn(room_id, player)
elif msg.kind == "say":
text = msg.payload.get("message", "").strip()
if not text: continue
if text.startswith("/"):
parts = text.split(); cmd = parts[0].lower()
if cmd == "/roll":
notation = parts[1] if len(parts) > 1 else "1d20"
result = game_manager.roll(notation)
if result:
roll_msg = room_manager.add_message(room_id, player.id, player.name, f"rolls {result.as_string}")
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="chat", payload={**roll_msg.model_dump(), "is_roll": True,}))
elif cmd == "/save":
room_manager.save_room_state(room_id)
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="system", payload={"message": f"Game progress saved by {player.name}."}))
elif cmd == "/remember":
memory_text = " ".join(parts[1:])
if memory_text:
story_manager.memory_manager.add_memory(room_id, memory_text)
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="system", payload={"message": f"{player.name} added a memory: '{memory_text[:50]}...'" }))
elif cmd == "/next":
await _advance_turn(room_id, player)
elif cmd == "/ooc":
ooc_text = " ".join(parts[1:])
if ooc_text:
ooc_msg = room_manager.add_message(room_id, player.id, player.name, f"// {ooc_text}")
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="chat", payload={**ooc_msg.model_dump(), "is_ooc": True, }))
else:
await websocket.send_text(WSOutgoingMessage(kind="system", payload={"message": f"Unknown command: {cmd}"}).model_dump_json())
else:
if room_state.turn_state == "GM_PROCESSING": continue
room_state.current_turn_actions[player.id] = text
action_msg = room_manager.add_message(room_id, player.id, player.name, text)
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="chat", payload=action_msg.model_dump()))
await connection_manager.broadcast(room_id, WSOutgoingMessage(kind="state_update", payload=room_state.model_dump()))
except Exception:
logger.error(f"Error processing message from {player.name}", exc_info=True)
except WebSocketDisconnect:
connection_manager.disconnect(room_id, websocket)
disconnected_player = room_manager.remove_player(room_id, player_id)
if disconnected_player and (room_state := room_manager.get_room(room_id)):
room_state.current_turn_actions.pop(player_id, None)
# If the disconnecting player was the host, clear the temporary host ID.
if room_state.host_player_id == player_id:
room_state.host_player_id = None
await connection_manager.broadcast(
room_id,
WSOutgoingMessage(kind="system", payload={"message": f"{disconnected_player.name} has left the game."}),
)
await connection_manager.broadcast(
room_id,
WSOutgoingMessage(kind="state_update", payload=room_state.model_dump()),
)