the-nexus/server.py

#!/usr/bin/env python3
"""
The Nexus — Unified HTTP + WebSocket server.

Serves static frontend files (Three.js app) over HTTP on port 8080
and runs the WebSocket gateway on port 8765.
Single-process, single-command deployment — no nginx required.
"""
import asyncio
import json
import logging
import mimetypes
import os
import signal
import sys
from http import HTTPStatus
from pathlib import Path
from typing import Set

import websockets
import websockets.asyncio.server

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
HTTP_PORT = int(os.environ.get("NEXUS_HTTP_PORT", "8080"))
WS_PORT = int(os.environ.get("NEXUS_WS_PORT", "8765"))
HOST = os.environ.get("NEXUS_HOST", "0.0.0.0")
ROOT = Path(__file__).resolve().parent

# Static file extensions we're willing to serve
SAFE_SUFFIXES = {
    ".html", ".htm", ".css", ".js", ".mjs", ".json",
    ".png", ".jpg", ".jpeg", ".gif", ".svg", ".ico",
    ".woff", ".woff2", ".ttf", ".eot",
    ".txt", ".xml", ".webmanifest",
}

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

# ---------------------------------------------------------------------------
# HTTP — static file server
# ---------------------------------------------------------------------------
async def http_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    """Minimal async HTTP/1.1 server for static files."""
    try:
        # Read request line + headers (max 8 KB)
        data = b""
        while b"\r\n\r\n" not in data and len(data) < 8192:
            chunk = await asyncio.wait_for(reader.read(4096), timeout=10)
            if not chunk:
                return
            data += chunk
        header_text = data.split(b"\r\n\r\n", 1)[0].decode("utf-8", errors="replace")
        lines = header_text.split("\r\n")
        request_line = lines[0]
        parts = request_line.split(" ", 2)
        if len(parts) < 2:
            writer.close()
            return
        method, raw_path = parts[0], parts[1]

        if method not in ("GET", "HEAD"):
            _write_response(writer, HTTPStatus.METHOD_NOT_ALLOWED, b"Method Not Allowed")
            return

        # Normalise path — prevent directory traversal
        safe_path = raw_path.split("?", 1)[0].split("#", 1)[0]
        safe_path = os.path.normpath(safe_path).lstrip("/")
        if not safe_path:
            safe_path = "index.html"

        file_path = ROOT / safe_path
        # Reject traversal
        if not str(file_path.resolve()).startswith(str(ROOT)):
            _write_response(writer, HTTPStatus.FORBIDDEN, b"Forbidden")
            return

        if not file_path.exists() or file_path.is_dir():
            # Try index.html for directories
            if file_path.is_dir() and (file_path / "index.html").exists():
                file_path = file_path / "index.html"
            else:
                _write_response(writer, HTTPStatus.NOT_FOUND, b"Not Found")
                return

        suffix = file_path.suffix.lower()
        if suffix not in SAFE_SUFFIXES:
            _write_response(writer, HTTPStatus.FORBIDDEN, b"Forbidden")
            return

        content_type = mimetypes.guess_type(str(file_path))[0] or "application/octet-stream"
        body = file_path.read_bytes()
        headers = {
            "Content-Type": content_type,
            "Content-Length": str(len(body)),
            "Cache-Control": "no-cache",
            "Access-Control-Allow-Origin": "*",
        }
        _write_response(writer, HTTPStatus.OK, body, headers, method == "HEAD")
    except (asyncio.TimeoutError, ConnectionError):
        pass
    except Exception as exc:
        log.warning("HTTP handler error: %s", exc)
    finally:
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass


def _write_response(writer, status, body, headers=None, head_only=False):
    line = f"HTTP/1.1 {status.value} {status.phrase}\r\n"
    hdr = ""
    if headers:
        for k, v in headers.items():
            hdr += f"{k}: {v}\r\n"
    response = (line + hdr + "\r\n").encode()
    if not head_only:
        response += body if isinstance(body, bytes) else body.encode()
    writer.write(response)


# ---------------------------------------------------------------------------
# WebSocket — broadcast gateway
# ---------------------------------------------------------------------------
ws_clients: Set[websockets.asyncio.server.ServerConnection] = set()


async def ws_handler(websocket: websockets.asyncio.server.ServerConnection):
    ws_clients.add(websocket)
    log.info("WS client connected from %s. Total: %d", websocket.remote_address, len(ws_clients))
    try:
        async for message in websocket:
            try:
                data = json.loads(message)
                msg_type = data.get("type", "unknown")
                if msg_type in ("agent_register", "thought", "action"):
                    log.debug("WS %s from %s", msg_type, websocket.remote_address)
            except (json.JSONDecodeError, TypeError):
                pass

            # Broadcast to all OTHER clients
            disconnected = set()
            for client in ws_clients:
                if client != websocket and client.open:
                    try:
                        await client.send(message)
                    except Exception:
                        disconnected.add(client)
            ws_clients.difference_update(disconnected)

    except websockets.exceptions.ConnectionClosed:
        pass
    except Exception as exc:
        log.error("WS handler error for %s: %s", websocket.remote_address, exc)
    finally:
        ws_clients.discard(websocket)
        log.info("WS client disconnected. Total: %d", len(ws_clients))


# ---------------------------------------------------------------------------
# Main — run both servers concurrently
# ---------------------------------------------------------------------------
async def main():
    # HTTP server
    http_server = await asyncio.start_server(http_handler, HOST, HTTP_PORT)
    log.info("HTTP server listening on http://%s:%d", HOST, HTTP_PORT)

    # WebSocket server
    ws_server = await websockets.asyncio.server.serve(ws_handler, HOST, WS_PORT)
    log.info("WebSocket server listening on ws://%s:%d", HOST, WS_PORT)
    log.info("Nexus is live — open http://%s:%d in a browser", HOST, HTTP_PORT)

    # 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:
            pass

    await stop
    log.info("Shutting down...")

    http_server.close()
    await http_server.wait_closed()
    ws_server.close()
    await ws_server.wait_closed()

    remaining = {c for c in ws_clients if c.open}
    if remaining:
        await asyncio.gather(*(c.close() for c in remaining), return_exceptions=True)
    ws_clients.clear()
    log.info("Shutdown complete.")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass
    except Exception as exc:
        log.critical("Fatal: %s", exc)
        sys.exit(1)