wu981526092's picture
Security: Fix critical vulnerabilities before public release
bcbd2ec
"""
FastAPI Application Definition
This module defines the FastAPI application and routes for the agent monitoring system.
"""
import logging
import os
import secrets
from pathlib import Path
import sys
from fastapi import FastAPI, Request, status, Depends, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from starlette.middleware.sessions import SessionMiddleware
from fastapi.responses import RedirectResponse, HTMLResponse, JSONResponse
from fastapi.exception_handlers import http_exception_handler
from backend.middleware.auth import ConditionalAuthMiddleware
from backend.middleware.usage_tracker import UsageTrackingMiddleware
from backend.dependencies import require_auth_in_hf_spaces
from utils.environment import should_enable_auth, debug_environment
# Add server module to path if not already there
server_dir = os.path.dirname(os.path.abspath(__file__))
if server_dir not in sys.path:
sys.path.append(server_dir)
# Import from backend modules
from backend.server_config import ensure_directories
from backend.routers import (
knowledge_graphs,
traces,
tasks,
temporal_graphs,
graph_comparison,
agentgraph,
example_traces,
methods,
observability,
auth,
testing,
)
# Setup logging
logger = logging.getLogger("agent_monitoring_server")
# Create FastAPI app
app = FastAPI(title="Agent Monitoring System", version="1.0.0")
# Define allowed CORS origins (security fix: restrict from wildcard)
ALLOWED_ORIGINS = [
"http://localhost:3001",
"http://localhost:5280",
"http://localhost:7860",
"https://holistic-ai-agentgraph.hf.space",
"https://huggingface.co",
]
# Add CORS middleware (first, so it's outermost)
app.add_middleware(
CORSMiddleware,
allow_origins=ALLOWED_ORIGINS,
allow_credentials=True,
allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
allow_headers=["Content-Type", "Authorization"],
)
# IMPORTANT: Middleware runs in REVERSE order of add_middleware() calls!
# We need Session to run BEFORE Auth so that request.session is available.
# Therefore: Add Auth FIRST, then Session SECOND.
# Detect if running in HF Spaces and adjust session config accordingly
from utils.environment import is_huggingface_space
# Debug environment detection
logger.info(f"πŸ” Environment detection: is_huggingface_space() = {is_huggingface_space()}")
logger.info(f"πŸ” SPACE_ID env var: {os.getenv('SPACE_ID')}")
# Add auth middleware FIRST (will run AFTER session middleware)
app.add_middleware(ConditionalAuthMiddleware)
# Add session middleware SECOND (will run BEFORE auth middleware, setting up request.session)
# HF Spaces specific session configuration
session_secret = os.getenv("SESSION_SECRET_KEY") or secrets.token_urlsafe(32)
if is_huggingface_space():
# HF Spaces optimized session configuration
logger.info("πŸ—οΈ Configuring session middleware for HF Spaces environment")
app.add_middleware(
SessionMiddleware,
secret_key=session_secret,
max_age=3600, # Shorter expiry for HF Spaces (1 hour)
same_site="none", # CRITICAL: Required for iframe cookies in HF Spaces
https_only=True, # HF Spaces uses HTTPS
)
else:
# Local development session configuration
logger.info("🏠 Configuring session middleware for local development")
app.add_middleware(
SessionMiddleware,
secret_key=session_secret,
max_age=86400, # 24 hours for local dev
same_site="lax", # Better for OAuth redirects
https_only=False, # HTTP for local dev
)
# Add usage tracking middleware (last added = runs first, outermost)
app.add_middleware(UsageTrackingMiddleware)
# Custom exception handler for authentication redirects
@app.exception_handler(HTTPException)
async def custom_http_exception_handler(request: Request, exc: HTTPException):
"""
Custom exception handler that redirects web requests to login page
when they receive 401 errors with redirect_to field.
"""
# Check if this is a 401 with redirect_to field
if exc.status_code == 401 and isinstance(exc.detail, dict) and "redirect_to" in exc.detail:
# Check if this is not an API request
is_api_request = (
request.url.path.startswith("/api/") or
request.headers.get("accept", "").startswith("application/json") or
request.headers.get("content-type", "").startswith("application/json")
)
if not is_api_request:
# Redirect web requests to login page
logger.info(f"πŸ”„ Redirecting unauthenticated user from {request.url.path} to login page")
return RedirectResponse(url=exc.detail["redirect_to"], status_code=302)
# For all other cases, use default exception handler
return await http_exception_handler(request, exc)
# Mount datasets directory for accessing json files
app.mount("/data", StaticFiles(directory="datasets"), name="data")
# Mount static directory for papers and other static assets
app.mount("/static", StaticFiles(directory="backend/static"), name="static")
# Include routers
app.include_router(auth.router) # Add auth router first
app.include_router(traces.router)
app.include_router(knowledge_graphs.router)
app.include_router(agentgraph.router)
app.include_router(tasks.router)
app.include_router(temporal_graphs.router)
app.include_router(graph_comparison.router)
app.include_router(example_traces.router)
app.include_router(methods.router)
app.include_router(observability.router)
app.include_router(testing.router)
# Start background scheduler for automated tasks
# scheduler_service.start()
@app.on_event("startup")
async def startup_event():
"""Start background services on app startup"""
logger.info("βœ… Backend server starting...")
# 🌍 Debug environment information
debug_environment()
# πŸ”§ Create necessary directories
ensure_directories()
logger.info("πŸ“ Directory structure created")
# πŸ—„οΈ Initialize database on startup
try:
from backend.database.init_db import init_database
init_database(reset=False, force=False)
logger.info("πŸ—„οΈ Database initialized successfully")
except Exception as e:
logger.error(f"❌ Database initialization failed: {e}")
# Don't fail startup - continue with empty database
# πŸ” Log authentication status
if should_enable_auth():
logger.info("πŸ” Authentication ENABLED (HF Spaces environment)")
else:
logger.info("🏠 Authentication DISABLED (Local development)")
logger.info("πŸš€ Backend API available at: http://0.0.0.0:7860")
# scheduler_service.start() # This line is now commented out
@app.on_event("shutdown")
async def shutdown_event():
"""Stop background services on app shutdown"""
logger.info("Stopping background services...")
# scheduler_service.stop() # This line is now commented out
# Root route - serve React app directly (authentication handled by frontend)
@app.get("/")
async def root(request: Request):
"""Serve the React app directly - authentication is now handled by frontend"""
return RedirectResponse(url="/agentgraph")
# Serve React app for any unmatched routes
@app.get("/app/{path:path}")
async def serve_react_app(path: str):
"""Serve the React app for client-side routing"""
return RedirectResponse(url="/agentgraph")