Update app/main.py

app/main.py

@@ -14,7 +14,7 @@ from app.config import get_settings
 from app.executor import get_executor
 from app.routes import health, mcp, tools
 
-# ─── Logging Configuration ────────────────────────────────────────
+# ─── Logging Configuration ────────────────────────────────────
 
 logging.basicConfig(
     level=logging.INFO,
@@ -25,7 +25,7 @@ logging.basicConfig(
 logger = logging.getLogger("mcp_code_executor")
 
 
-# ─── Application Lifespan ─────────────────────────────────────────
+# ─── Application Lifespan ─────────────────────────────────────
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -38,6 +38,7 @@ async def lifespan(app: FastAPI):
     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
@@ -49,13 +50,14 @@ async def lifespan(app: FastAPI):
     logger.info("Shutdown complete")
 
 
-# ─── FastAPI Application ──────────────────────────────────────────
+# ─── 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 incremental code generation, "
+        "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",
@@ -65,7 +67,7 @@ app = FastAPI(
     openapi_url="/openapi.json",
 )
 
-# ─── CORS Middleware ───────────────────────────────────────────────
+# ─── CORS Middleware ───────────────────────────────────────────
 
 app.add_middleware(
     CORSMiddleware,
@@ -73,10 +75,26 @@ app.add_middleware(
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
+    expose_headers=["*"],
 )
 
 
-# ─── Global Exception Handler ─────────────────────────────────────
+# ─── 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):
@@ -90,14 +108,14 @@ async def global_exception_handler(request: Request, exc: Exception):
     )
 
 
-# ─── Register Routes ──────────────────────────────────────────────
+# ─── Register Routes ──────────────────────────────────────────
 
 app.include_router(health.router)
 app.include_router(mcp.router)
 app.include_router(tools.router)
 
 
-# ─── Direct Run (development) ─────────────────────────────────────
+# ─── Direct Run (development) ─────────────────────────────────
 
 if __name__ == "__main__":
     import uvicorn