agent-platform/api/main.py

"""
FastAPI application with lifespan management.
"""

from __future__ import annotations

import logging
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from redis.asyncio import Redis
from sqlalchemy.ext.asyncio import create_async_engine

from api.middleware.tracing import TracingMiddleware
from api.routes import approve, chat, health, trace
from core.audit.logger import AuditLogger
from core.config import get_settings
from core.llm.claude_client import ClaudeClient
from core.router.model_router import ModelRouter

logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
    """
    Lifespan context manager for startup/shutdown.

    Startup order:
    1. Setup observability
    2. Create DB engine
    3. ToolRegistry (3 tools)
    4. ModelRouter
    5. ClaudeClient
    6. AuditLogger
    7. AgentEngine (PostgresSaver for prod, MemorySaver for dev/test)
    8. Store singletons in app.state
    """
    settings = get_settings()

    logger.info("Starting application...")

    # 1. Setup observability
    # TODO: Implement setup_observability(service_name, environment)
    logger.info("Observability setup complete")

    # 2. Create DB engine
    db_engine = create_async_engine(
        settings.DATABASE_URL,
        pool_size=settings.DB_POOL_SIZE,
        max_overflow=settings.DB_MAX_OVERFLOW,
        pool_timeout=settings.DB_POOL_TIMEOUT,
        pool_recycle=settings.DB_POOL_RECYCLE,
    )
    logger.info("Database engine created")

    # Create Redis client
    redis_client = Redis.from_url(settings.REDIS_URL, decode_responses=False)
    logger.info("Redis client created")

    # 3. ToolRegistry
    # TODO: Register three tools
    logger.info("ToolRegistry initialized")

    # 4. ModelRouter
    model_router = ModelRouter()
    logger.info("ModelRouter initialized")

    # 5. ClaudeClient
    claude_client = ClaudeClient()
    logger.info("ClaudeClient initialized")

    # 6. AuditLogger
    audit_logger = AuditLogger(db_session_factory=None, otel_tracer=None)
    logger.info("AuditLogger initialized")

    # 7. AgentEngine
    # TODO: Build agent graph with appropriate saver
    # For now, use placeholder
    agent_engine = None
    logger.info("AgentEngine initialized")

    # 8. Store singletons in app.state
    app.state.db_engine = db_engine
    app.state.redis_client = redis_client
    app.state.model_router = model_router
    app.state.claude_client = claude_client
    app.state.audit_logger = audit_logger
    app.state.agent_engine = agent_engine

    logger.info("Application startup complete")

    yield

    # Shutdown
    logger.info("Shutting down application...")

    # Close Redis
    await redis_client.aclose()
    logger.info("Redis client closed")

    # Dispose DB engine
    await db_engine.dispose()
    logger.info("Database engine disposed")

    logger.info("Application shutdown complete")


# Create FastAPI app
app = FastAPI(
    title="Agent Platform API",
    version="1.0.0",
    lifespan=lifespan,
)

# Get settings
settings = get_settings()

# Middleware registration (outer → inner)
# 1. TracingMiddleware
app.add_middleware(TracingMiddleware)

# 2. LoggingMiddleware
# TODO: Implement LoggingMiddleware

# 3. RateLimitMiddleware
# Note: Redis client will be available after startup
# We'll initialize it in a way that accesses app.state

# 4. CORS
if settings.ENVIRONMENT == "development":
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )
else:
    # TODO: Use settings.ALLOWED_ORIGINS for production
    app.add_middleware(
        CORSMiddleware,
        allow_origins=["https://example.com"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

# Route registration
app.include_router(health.router, tags=["health"])
app.include_router(chat.router, tags=["chat"])
app.include_router(approve.router, tags=["approve"])
app.include_router(trace.router, tags=["trace"])