mirror of
https://github.com/Nighthawk42/VDM.git
synced 2026-08-30 04:30:21 +00:00
548 lines
21 KiB
Python
548 lines
21 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 ---
|
|
# UPDATED: Initialize the DatabaseManager with separate paths for sessions and users.
|
|
db_manager = DatabaseManager(
|
|
sessions_db_path=Path(settings.memory.sessions_db_file),
|
|
users_db_path=Path(settings.memory.users_db_file)
|
|
)
|
|
|
|
# Managers that depend on the database manager
|
|
persistence_manager = PersistenceManager(db_manager=db_manager)
|
|
user_manager = UserManager(db_manager=db_manager)
|
|
|
|
# RoomManager depends on other managers
|
|
room_manager = RoomManager(user_manager=user_manager, persistence_manager=persistence_manager)
|
|
|
|
# Standalone managers
|
|
story_manager = StoryManager()
|
|
audio_manager = AudioManager()
|
|
game_manager = DiceRoller()
|
|
connection_manager = ConnectionManager()
|
|
|
|
|
|
# ===================================================================
|
|
# Core Game Loop Logic
|
|
# ===================================================================
|
|
|
|
# ... (the rest of the file is unchanged) ...
|
|
|
|
|
|
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) # Pass room_id for audio path
|
|
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) # Pass room_id for audio path
|
|
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) # Pass room_id for audio path
|
|
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) # Pass room_id for audio path
|
|
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)
|
|
if (
|
|
not room_state
|
|
or player.id != room_state.host_player_id
|
|
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) # Pass room_id for audio path
|
|
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.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
|
|
|
|
# Send the existing chat history to the newly connected player
|
|
if room.messages:
|
|
await websocket.send_text(
|
|
WSOutgoingMessage(
|
|
kind="chat_history", payload={"messages": [m.model_dump() for m in room.messages]}
|
|
).model_dump_json()
|
|
)
|
|
|
|
|
|
if not room.host_player_id:
|
|
room.host_player_id = player_id
|
|
logger.info(f"Player '{player.name}' is now the host of room '{room_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
|
|
|
|
if msg.kind == "start_game":
|
|
if (
|
|
player.id == room_state.host_player_id
|
|
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":
|
|
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
|
|
|
|
is_command = text.startswith("/")
|
|
if is_command:
|
|
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:
|
|
user_manager.logout(player_token)
|
|
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 room_state.host_player_id == player_id:
|
|
new_host_id = next(
|
|
(pid for pid, p in room_state.players.items() if p.is_active), None
|
|
)
|
|
room_state.host_player_id = new_host_id
|
|
if new_host_id:
|
|
new_host_name = room_state.players[new_host_id].name
|
|
await connection_manager.broadcast(
|
|
room_id,
|
|
WSOutgoingMessage(
|
|
kind="system",
|
|
payload={
|
|
"message": f"The host has left. {new_host_name} is the new host."
|
|
},
|
|
),
|
|
)
|
|
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()),
|
|
) |