app/main.py

"""FastAPI application entry point - MCP Code Executor."""

from __future__ import annotations

import logging
import sys
from contextlib import asynccontextmanager

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

from app.config import get_settings
from app.executor import get_executor
from app.routes import health, mcp, tools

# ─── Logging Configuration ────────────────────────────────────

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    stream=sys.stdout,
)
logger = logging.getLogger("mcp_code_executor")


# ─── Application Lifespan ─────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup and shutdown logic."""
    settings = get_settings()
    logger.info("=" * 60)
    logger.info("MCP Code Executor starting up")
    logger.info("  Environment: %s", settings.env_type.value)
    logger.info("  Code storage: %s", settings.code_storage_dir)
    logger.info("  Max concurrent: %d", settings.max_concurrent_executions)
    logger.info("  Execution timeout: %ds", settings.execution_timeout)
    logger.info("  Python executable: %s", settings.get_python_executable())
    logger.info("  MCP endpoint: /mcp/jsonrpc")
    logger.info("=" * 60)

    yield

    # Shutdown: cleanup active processes
    logger.info("Shutting down - cleaning up active processes...")
    executor = get_executor()
    await executor.cleanup()
    logger.info("Shutdown complete")


# ─── FastAPI Application ──────────────────────────────────────

app = FastAPI(
    title="MCP Code Executor",
    description=(
        "An MCP server that allows LLMs to execute Python code within a "
        "specified Python environment. Supports the full MCP 2025-03-26 "
        "Streamable HTTP transport protocol, incremental code generation, "
        "package management, and concurrent execution."
    ),
    version="1.0.0",
    lifespan=lifespan,
    docs_url="/docs",
    redoc_url="/redoc",
    openapi_url="/openapi.json",
)

# ─── CORS Middleware ───────────────────────────────────────────

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
)


# ─── MCP Headers Middleware ────────────────────────────────────

@app.middleware("http")
async def add_mcp_headers(request: Request, call_next):
    """Add MCP-required headers to responses on MCP endpoints."""
    response = await call_next(request)

    if request.url.path.startswith("/mcp/"):
        # MCP Streamable HTTP requires these headers
        response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
        response.headers["Connection"] = "keep-alive"

    return response


# ─── Global Exception Handler ─────────────────────────────────

@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    logger.exception("Unhandled exception on %s %s", request.method, request.url.path)
    return JSONResponse(
        status_code=500,
        content={
            "error": "Internal server error",
            "detail": str(exc),
        },
    )


# ─── Register Routes ──────────────────────────────────────────

app.include_router(health.router)
app.include_router(mcp.router)
app.include_router(tools.router)


# ─── Direct Run (development) ─────────────────────────────────

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=7860,
        reload=True,
        log_level="info",
    )