iLSP/ilsp/server.py

"""
iLSP main entrypoint.

Starts all services in one container:
  - pylsp        on localhost:2087  (internal only — Python LSP + i80 completions)
  - Bicep proxy  on localhost:2088  (internal only — Bicep LS wrapper + LRU modules)
  - aiohttp HTTP on $HTTP_PORT      (exposed — health, reload, WebSocket proxies)

Only the HTTP port is exposed to Traefik/Nomad. Editors connect via:
  ws://ilsp.i80.dk/python   →  pylsp
  ws://ilsp.i80.dk/bicep    →  Bicep LS
"""

import asyncio
import logging
import os
import signal
import subprocess
import sys
import threading

from aiohttp import web

from .python_lsp.catalog import PypiCatalog
from .bicep_lsp.modules import BicepModuleCatalog
from .bicep_lsp.proxy import serve_bicep

logger = logging.getLogger(__name__)

PYTHON_LSP_PORT = int(os.getenv("PYTHON_LSP_PORT", "2087"))
BICEP_LSP_PORT  = int(os.getenv("BICEP_LSP_PORT",  "2088"))
HTTP_PORT       = int(os.getenv("HTTP_PORT", "8000"))

_CHUNK = 65536


def _serve_python_lsp(port: int) -> None:
    """Start pylsp in TCP mode on localhost only; restart on unexpected exit."""
    while True:
        proc = subprocess.Popen(
            [sys.executable, "-m", "pylsp", "--tcp", "--host", "127.0.0.1", "--port", str(port)],
        )
        logger.info("pylsp listening on localhost:%d  PID=%d", port, proc.pid)
        proc.wait()
        logger.warning("pylsp exited (code %s) — restarting", proc.returncode)


async def _pipe(reader: asyncio.StreamReader, writer) -> None:
    """Pipe bytes from reader to writer until EOF."""
    try:
        while True:
            data = await reader.read(_CHUNK)
            if not data:
                break
            writer.write(data)
            if hasattr(writer, "drain"):
                await writer.drain()
            else:
                await writer.send_bytes(data)
    except Exception:
        pass


async def _ws_proxy(request: web.Request, host: str, port: int) -> web.WebSocketResponse:
    """Bridge an incoming WebSocket connection to an internal TCP LSP server."""
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    try:
        tcp_reader, tcp_writer = await asyncio.open_connection(host, port)
    except OSError as exc:
        logger.error("Cannot connect to LSP on %s:%d — %s", host, port, exc)
        await ws.close(code=1011, message=b"LSP backend unavailable")
        return ws

    async def ws_to_tcp():
        try:
            async for msg in ws:
                if msg.type == web.WSMsgType.BINARY:
                    tcp_writer.write(msg.data)
                    await tcp_writer.drain()
                elif msg.type == web.WSMsgType.TEXT:
                    tcp_writer.write(msg.data.encode())
                    await tcp_writer.drain()
        except Exception:
            pass
        finally:
            tcp_writer.close()

    async def tcp_to_ws():
        try:
            while True:
                data = await tcp_reader.read(_CHUNK)
                if not data:
                    break
                await ws.send_bytes(data)
        except Exception:
            pass
        finally:
            await ws.close()

    await asyncio.gather(ws_to_tcp(), tcp_to_ws(), return_exceptions=True)
    return ws


async def _build_app() -> web.Application:
    app = web.Application()

    async def health(_: web.Request) -> web.Response:
        return web.json_response({
            "status": "ok",
            "pypi_packages": len(PypiCatalog._packages),
            "bicep_modules": len(BicepModuleCatalog._modules),
        })

    async def reload(_: web.Request) -> web.Response:
        before = len(BicepModuleCatalog._modules)
        BicepModuleCatalog.load()
        after = len(BicepModuleCatalog._modules)
        logger.info("Catalog reloaded: %d → %d modules", before, after)
        return web.json_response({
            "status": "reloaded",
            "bicep_modules_before": before,
            "bicep_modules_after": after,
        })

    async def python_ws(request: web.Request) -> web.WebSocketResponse:
        return await _ws_proxy(request, "127.0.0.1", PYTHON_LSP_PORT)

    async def bicep_ws(request: web.Request) -> web.WebSocketResponse:
        return await _ws_proxy(request, "127.0.0.1", BICEP_LSP_PORT)

    app.router.add_get("/health", health)
    app.router.add_post("/reload", reload)
    app.router.add_get("/python", python_ws)
    app.router.add_get("/bicep", bicep_ws)
    return app


async def main_async() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    logger.info("Pre-warming catalogs…")
    BicepModuleCatalog.load()
    await PypiCatalog.start_background_refresh()

    # LSP servers run internally on localhost — not exposed outside the container
    threading.Thread(target=_serve_python_lsp, args=(PYTHON_LSP_PORT,), daemon=True).start()
    threading.Thread(target=serve_bicep, args=(BICEP_LSP_PORT,), daemon=True).start()

    app = await _build_app()
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "0.0.0.0", HTTP_PORT)
    await site.start()
    logger.info("iLSP HTTP on http://0.0.0.0:%d  (ws://.../python  ws://.../bicep)", HTTP_PORT)

    await asyncio.Event().wait()


def main() -> None:
    loop = asyncio.new_event_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, loop.stop)
    loop.run_until_complete(main_async())