the-nexus/server.py

#!/usr/bin/env python3
"""
The Nexus WebSocket Gateway — Robust broadcast bridge for Timmy's consciousness.
This server acts as the central hub for the-nexus, connecting the mind (nexus_think.py),
the body (Evennia/Morrowind), and the visualization surface.
"""
import asyncio
import json
import logging
import os
import signal
import sys
from typing import Set

# Branch protected file - see POLICY.md
import websockets

# Configuration
PORT = int(os.environ.get("NEXUS_WS_PORT", "8765"))
HOST = os.environ.get("NEXUS_WS_HOST", "127.0.0.1")  # Local-only by default
# Set NEXUS_WS_HOST=0.0.0.0 to allow external connections (requires authentication)

# Logging setup
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger("nexus-gateway")

# State
clients: Set[websockets.WebSocketServerProtocol] = set()

async def broadcast_handler(websocket: websockets.WebSocketServerProtocol):
    """Handles individual client connections and message broadcasting."""

    # Authentication check for external connections
    if HOST != "127.0.0.1":
        # Require authentication token for external connections
        auth_token = os.environ.get("NEXUS_WS_AUTH_TOKEN")
        if not auth_token:
            logger.warning("External connections require NEXUS_WS_AUTH_TOKEN to be set")
            await websocket.close(1008, "Authentication required")
            return
        
        # Check for authentication in first message
        try:
            auth_message = await asyncio.wait_for(websocket.recv(), timeout=5.0)
            auth_data = json.loads(auth_message)
            
            if auth_data.get("type") != "auth" or auth_data.get("token") != auth_token:
                logger.warning(f"Authentication failed from {websocket.remote_address}")
                await websocket.close(1008, "Authentication failed")
                return
                
            logger.info(f"Authenticated connection from {websocket.remote_address}")
            
        except (asyncio.TimeoutError, json.JSONDecodeError, Exception) as e:
            logger.warning(f"Authentication error from {websocket.remote_address}: {e}")
            await websocket.close(1008, "Authentication required")
            return

    clients.add(websocket)
    addr = websocket.remote_address
    logger.info(f"Client connected from {addr}. Total clients: {len(clients)}")
    
    try:
        async for message in websocket:
            # Parse for logging/validation if it's JSON
            try:
                data = json.loads(message)
                msg_type = data.get("type", "unknown")
                # Optional: log specific important message types
                if msg_type in ["agent_register", "thought", "action"]:
                    logger.debug(f"Received {msg_type} from {addr}")
            except (json.JSONDecodeError, TypeError):
                pass

            # Broadcast to all OTHER clients
            if not clients:
                continue
                
            disconnected = set()
            # Create broadcast tasks, tracking which client each task targets
            task_client_pairs = []
            for client in clients:
                if client != websocket and client.open:
                    task = asyncio.create_task(client.send(message))
                    task_client_pairs.append((task, client))

            if task_client_pairs:
                tasks = [pair[0] for pair in task_client_pairs]
                results = await asyncio.gather(*tasks, return_exceptions=True)
                for i, result in enumerate(results):
                    if isinstance(result, Exception):
                        target_client = task_client_pairs[i][1]
                        logger.error(f"Failed to send to client {target_client.remote_address}: {result}")
                        disconnected.add(target_client)
            
            if disconnected:
                clients.difference_update(disconnected)
                
    except websockets.exceptions.ConnectionClosed:
        logger.debug(f"Connection closed by client {addr}")
    except Exception as e:
        logger.error(f"Error handling client {addr}: {e}")
    finally:
        clients.discard(websocket)
        logger.info(f"Client disconnected {addr}. Total clients: {len(clients)}")

async def main():
    """Main server loop with graceful shutdown."""
    logger.info(f"Starting Nexus WS gateway on ws://{HOST}:{PORT}")
    
    # Set up signal handlers for graceful shutdown
    loop = asyncio.get_running_loop()
    stop = loop.create_future()
    
    def shutdown():
        if not stop.done():
            stop.set_result(None)

    for sig in (signal.SIGINT, signal.SIGTERM):
        try:
            loop.add_signal_handler(sig, shutdown)
        except NotImplementedError:
            # Signal handlers not supported on Windows
            pass

    async with websockets.serve(broadcast_handler, HOST, PORT):
        logger.info("Gateway is ready and listening.")
        await stop
        
    logger.info("Shutting down Nexus WS gateway...")
    # Close any remaining client connections (handlers may have already cleaned up)
    remaining = {c for c in clients if c.open}
    if remaining:
        logger.info(f"Closing {len(remaining)} active connections...")
        close_tasks = [client.close() for client in remaining]
        await asyncio.gather(*close_tasks, return_exceptions=True)
    clients.clear()
    
    logger.info("Shutdown complete.")

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass
    except Exception as e:
        logger.critical(f"Fatal server error: {e}")
        sys.exit(1)