http_server.py

#!/usr/bin/env python3
"""
Persistent HTTP server for code-rag MCP system.

Runs as a long-lived process serving MCP via StreamableHTTP.
Projects are identified by the X-Project-Root header in each request.
The DB for each project lives at {project_root}/.code-rag/milvus.db.

Start: ./code-rag-server.sh
Health: curl http://127.0.0.1:7101/health
"""

import contextlib
import json
import logging
import os
import signal
import sys
import traceback
from pathlib import Path

import uvicorn
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Mount, Route

from mcp.server import Server
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager

import rag_milvus
import file_watcher
from tools import register_tools, set_current_project_root

logger = logging.getLogger("code-rag")

# --- Configuration ---

HOST = os.getenv("CODE_RAG_HOST", "127.0.0.1")
PORT = int(os.getenv("CODE_RAG_PORT", "7101"))

# Server runtime files live in ~/.code-rag/ (not in any project)
_SERVER_DIR = Path.home() / ".code-rag"
PID_FILE = _SERVER_DIR / "server.pid"
LOG_FILE = _SERVER_DIR / "server.log"


# --- Project middleware ---

class ProjectMiddleware:
    """ASGI middleware that extracts X-Project-Root header and sets ContextVar."""

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        if scope["type"] == "http":
            headers = dict(scope.get("headers", []))
            project_root = headers.get(b"x-project-root", b"").decode("utf-8").strip()

            if project_root:
                set_current_project_root(project_root)
            else:
                set_current_project_root(None)

        try:
            await self.app(scope, receive, send)
        except Exception as exc:
            logger.error("ASGI handler error: %s\n%s", exc, traceback.format_exc())
            if scope["type"] == "http":
                body = json.dumps({"error": "internal_server_error", "detail": str(exc)}).encode()
                await send({"type": "http.response.start", "status": 500, "headers": [
                    [b"content-type", b"application/json"],
                    [b"content-length", str(len(body)).encode()],
                ]})
                await send({"type": "http.response.body", "body": body})


# --- Health endpoint ---

_model_loaded = False
_server_mode_ready = False


async def health(request: Request) -> JSONResponse:
    watcher_status = file_watcher.get_watcher_status()
    watchers = {}
    for root, s in watcher_status.items():
        watchers[Path(root).name] = {
            "pending": s["pending"],
            "processing": s["processing"],
            "files_indexed": s["stats"]["files_indexed"],
            "files_deleted": s["stats"]["files_deleted"],
            "batches": s["stats"]["batches_processed"],
        }
    return JSONResponse({
        "status": "ok",
        "model": _model_loaded,
        "embed_model": Path(rag_milvus._MODEL_PATH).name if _model_loaded else None,
        "embed_dim": rag_milvus._EMBED_DIM,
        "milvus": _server_mode_ready,
        "watchers": watchers,
    })


# --- Lifespan ---

@contextlib.asynccontextmanager
async def lifespan(app: Starlette):
    """Server lifecycle: PID file, model preload, server mode init."""
    _SERVER_DIR.mkdir(parents=True, exist_ok=True)

    # Write PID file
    PID_FILE.write_text(str(os.getpid()))
    print(f"[HTTP] PID {os.getpid()} written to {PID_FILE}", file=sys.stderr)

    global _model_loaded, _server_mode_ready

    # Pre-load MLX model
    try:
        print("[HTTP] Pre-loading MLX model...", file=sys.stderr)
        rag_milvus.get_mlx_model()
        _model_loaded = True
        print("[HTTP] MLX model loaded.", file=sys.stderr)
    except Exception as e:
        print(f"[HTTP] Warning: Could not pre-load model: {e}", file=sys.stderr)

    # Init server mode (lazy persistent clients per project)
    try:
        rag_milvus.init_server_mode()
        _server_mode_ready = True
    except Exception as e:
        print(f"[HTTP] Warning: Could not init server mode: {e}", file=sys.stderr)

    # Start session manager
    async with session_manager.run():
        print(f"[HTTP] Server ready on http://{HOST}:{PORT}", file=sys.stderr)
        try:
            yield
        finally:
            # Stop file watchers before closing Milvus clients
            await file_watcher.stop_all_watchers()

    # Cleanup
    rag_milvus.close_server_mode()
    if PID_FILE.exists():
        PID_FILE.unlink()
    print("[HTTP] Server stopped.", file=sys.stderr)


# --- MCP server setup ---

mcp_server = Server("code-rag")
register_tools(mcp_server)

session_manager = StreamableHTTPSessionManager(
    app=mcp_server,
    stateless=True,
    json_response=True,
)


# --- Starlette app ---

app = Starlette(
    routes=[
        Route("/health", health, methods=["GET"]),
        Mount("/mcp", app=ProjectMiddleware(session_manager.handle_request)),
    ],
    lifespan=lifespan,
)


# --- Signal handling ---

def _handle_signal(signum, frame):
    """Graceful shutdown on SIGTERM/SIGINT."""
    print(f"[HTTP] Received signal {signum}, shutting down...", file=sys.stderr)
    raise SystemExit(0)


if __name__ == "__main__":
    signal.signal(signal.SIGTERM, _handle_signal)

    uvicorn.run(
        app,
        host=HOST,
        port=PORT,
        log_level="warning",
    )