middleware.py•2.27 kB
"""
Middleware setup for MCPilot
"""
import time
import logging
from typing import Callable
from fastapi import FastAPI, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.trustedhost import TrustedHostMiddleware
logger = logging.getLogger(__name__)
def setup_middleware(app: FastAPI) -> None:
"""Setup all middleware for the FastAPI application"""
# CORS middleware
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Configure this based on settings
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Trusted hosts middleware (configure based on settings)
app.add_middleware(
TrustedHostMiddleware,
allowed_hosts=["*"] # Configure this based on settings
)
# Custom request logging middleware
@app.middleware("http")
async def log_requests(request: Request, call_next: Callable) -> Response:
"""Log all HTTP requests"""
start_time = time.time()
# Log request
logger.info(f"Request: {request.method} {request.url}")
# Process request
response = await call_next(request)
# Calculate processing time
process_time = time.time() - start_time
# Log response
logger.info(
f"Response: {response.status_code} - "
f"Processed in {process_time:.4f}s"
)
# Add processing time header
response.headers["X-Process-Time"] = str(process_time)
return response
# Custom error handling middleware
@app.middleware("http")
async def error_handling(request: Request, call_next: Callable) -> Response:
"""Handle errors and add consistent error responses"""
try:
response = await call_next(request)
return response
except Exception as exc:
logger.exception(f"Unhandled error in request {request.url}")
# Return a generic error response
from fastapi.responses import JSONResponse
return JSONResponse(
status_code=500,
content={"detail": "Internal server error"}
)