core/engine/src/main.py - main - alexefan136/flowstack

core/engine/src/main.py
748 строк23 KB
Alexander Efanov
upd fix
04 авг 2026, 14:09
04 авг 2026, 14:09f63ab86
"""
Главный модуль — FastAPI приложение для FlowStack Engine.

Lifecycle:
- Startup: init_db, reset stuck states, tools manager, model discovery
- Shutdown: model discovery, tools manager, MCP manager, RAG client, DB

Flows управляются через /api/v1/flows (DB-based, FlowService).
LLM endpoints — через /api/v1/llm (src.llm.routes).
"""

from __future__ import annotations

import logging
import time
import uuid
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from datetime import UTC, datetime
from typing import Any

import structlog
import uvicorn
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from sqlalchemy import select
from sqlalchemy import update as sql_update

from src import __version__
from src.api import (
    API_TAGS,
    get_api_info,
    register_all_routers,
)
from src.config import get_settings
from src.db.models import Agent, MCPServer, Task
from src.db.session import (
    async_session_maker,
    check_database_health,
    close_db,
    get_engine_stats,
    init_db,
)
from src.llm import (
    get_registry_stats,
    start_background_refresh,
    stop_background_refresh,
)
from src.llm import (
    router as llm_router,  # ✅ FIX: LLM-роутер для регистрации
)
from src.middleware.auth import AuthMiddleware
from src.runtime import shutdown_mcp_manager
from src.services import close_rag_client
from src.tools.registry import ToolRegistryManager

settings = get_settings()


# ============================================================================
# LOGGING CONFIGURATION
# ============================================================================


def configure_logging() -> None:
    """
    Настроить structlog.

    - Development: цветной console output
    - Production: JSON логи (для агрегации)
    """
    processors: list[Any] = [
        structlog.contextvars.merge_contextvars,
        structlog.processors.add_log_level,
        structlog.processors.TimeStamper(fmt="iso"),
        structlog.processors.StackInfoRenderer(),
        structlog.processors.format_exc_info,
    ]

    if settings.debug:
        processors.append(structlog.dev.ConsoleRenderer())
    else:
        processors.append(structlog.processors.JSONRenderer())

    structlog.configure(
        processors=processors,
        wrapper_class=structlog.make_filtering_bound_logger(
            logging.getLevelName(settings.log_level)
        ),
        logger_factory=structlog.PrintLoggerFactory(),
        cache_logger_on_first_use=True,
    )


configure_logging()
logger = structlog.get_logger()

_start_time: float = 0.0
_tools_manager: ToolRegistryManager | None = None


# ============================================================================
# Startup Helpers (reset stuck states)
# ============================================================================


async def reset_running_agents() -> None:
    """Сбросить 'зависшие' статусы агентов при старте."""
    try:
        async with async_session_maker() as session:
            stmt = select(Agent).where(Agent.status.in_(["running", "starting"]))
            result = await session.execute(stmt)
            stuck_agents = list(result.scalars().all())

            if not stuck_agents:
                return

            logger.warning(
                "engine.agents.reset_stuck",
                count=len(stuck_agents),
                agent_ids=[str(a.id) for a in stuck_agents],
            )

            update_stmt = (
                sql_update(Agent)
                .where(Agent.status.in_(["running", "starting"]))
                .values(
                    status="idle",
                    current_task_id=None,
                    updated_at=datetime.now(UTC),
                )
            )
            await session.execute(update_stmt)
            await session.commit()

            logger.info("engine.agents.reset_completed", count=len(stuck_agents))
    except Exception as e:
        logger.error("engine.agents.reset_failed", error=str(e), error_type=type(e).__name__)


async def reset_running_mcp_servers() -> None:
    """Сбросить 'зависшие' статусы MCP серверов при старте."""
    try:
        async with async_session_maker() as session:
            stmt = select(MCPServer).where(MCPServer.status.in_(["running", "starting"]))
            result = await session.execute(stmt)
            stuck_servers = list(result.scalars().all())

            if not stuck_servers:
                return

            logger.warning(
                "engine.mcp_servers.reset_stuck",
                count=len(stuck_servers),
                server_ids=[str(s.id) for s in stuck_servers],
            )

            update_stmt = (
                sql_update(MCPServer)
                .where(MCPServer.status.in_(["running", "starting"]))
                .values(status="stopped", updated_at=datetime.now(UTC))
            )
            await session.execute(update_stmt)
            await session.commit()

            logger.info("engine.mcp_servers.reset_completed", count=len(stuck_servers))
    except Exception as e:
        logger.error("engine.mcp_servers.reset_failed", error=str(e), error_type=type(e).__name__)


async def reset_running_tasks() -> None:
    """Сбросить задачи со статусом 'in_progress' при старте."""
    try:
        async with async_session_maker() as session:
            stmt = select(Task).where(Task.status == "in_progress")
            result = await session.execute(stmt)
            stuck_tasks = list(result.scalars().all())

            if not stuck_tasks:
                return

            logger.warning(
                "engine.tasks.reset_stuck",
                count=len(stuck_tasks),
                task_ids=[str(t.id) for t in stuck_tasks],
            )

            update_stmt = (
                sql_update(Task)
                .where(Task.status == "in_progress")
                .values(status="todo", updated_at=datetime.now(UTC))
            )
            await session.execute(update_stmt)
            await session.commit()

            logger.info("engine.tasks.reset_completed", count=len(stuck_tasks))
    except Exception as e:
        logger.error("engine.tasks.reset_failed", error=str(e), error_type=type(e).__name__)


async def _reset_all_stuck_states() -> None:
    """Запустить все сбросы зависших состояний."""
    await reset_running_agents()
    await reset_running_mcp_servers()
    await reset_running_tasks()


# ============================================================================
# Sentry & Prometheus Setup
# ============================================================================


def _setup_sentry() -> None:
    """Инициализация Sentry для error tracking."""
    sentry_dsn = getattr(settings, "sentry_dsn", None)
    if not sentry_dsn:
        return

    try:
        import sentry_sdk  # type: ignore[import-not-found,import-untyped]
        from sentry_sdk.integrations.fastapi import (
            FastApiIntegration,  # type: ignore[import-not-found,import-untyped]
        )
        from sentry_sdk.integrations.sqlalchemy import (
            SqlalchemyIntegration,  # type: ignore[import-not-found,import-untyped]
        )
        from sentry_sdk.integrations.starlette import (
            StarletteIntegration,  # type: ignore[import-not-found,import-untyped]
        )

        dsn_value = (
            sentry_dsn.get_secret_value()
            if hasattr(sentry_dsn, "get_secret_value")
            else str(sentry_dsn)
        )

        sentry_sdk.init(
            dsn=dsn_value,
            environment=settings.app_env,
            release=f"flowstack-engine@{__version__}",
            traces_sample_rate=0.1 if settings.app_env == "production" else 1.0,
            integrations=[
                StarletteIntegration(transaction_style="endpoint"),
                FastApiIntegration(transaction_style="endpoint"),
                SqlalchemyIntegration(),
            ],
        )
        logger.info("sentry.initialized", environment=settings.app_env)
    except ImportError:
        logger.warning("sentry.not_installed", hint="pip install sentry-sdk[fastapi]")
    except Exception as e:
        logger.error("sentry.init_failed", error=str(e))


def _setup_prometheus(app: FastAPI) -> None:
    """Инициализация Prometheus metrics."""
    if not getattr(settings, "enable_metrics", False):
        return

    try:
        from prometheus_fastapi_instrumentator import (
            Instrumentator,  # type: ignore[import-not-found,import-untyped]
        )

        Instrumentator(
            should_group_status_codes=True,
            should_ignore_untemplated=True,
            should_respect_env_var=False,
            excluded_handlers=["/metrics", "/health"],
        ).instrument(app).expose(
            app, endpoint="/metrics", include_in_schema=False, should_gzip=True
        )
        logger.info("prometheus.initialized", endpoint="/metrics")
    except ImportError:
        logger.warning(
            "prometheus.not_installed",
            hint="pip install prometheus-fastapi-instrumentator",
        )
    except Exception as e:
        logger.error("prometheus.init_failed", error=str(e))


# ============================================================================
# Lifespan
# ============================================================================


@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
    """Application lifespan manager."""
    global _start_time, _tools_manager
    _start_time = time.time()

    logger.info(
        "engine.starting",
        version=__version__,
        env=settings.app_env,
        debug=settings.debug,
    )

    _setup_sentry()

    # === Database ===
    try:
        await init_db()
        logger.info("engine.database.initialized")
    except Exception as e:
        logger.error("engine.database.failed", error=str(e))
        raise

    # === Reset stuck states ===
    await _reset_all_stuck_states()

    # === Tools Registry ===
    try:
        _tools_manager = ToolRegistryManager()
        _tools_manager.register_builtin_tools(
            file_ops_base_dir=str(getattr(settings, "file_ops_base_dir", "/tmp/flowstack-sandbox")),
            file_ops_read_only=bool(getattr(settings, "file_ops_read_only", False)),
        )
        _tools_manager.register_default_mcp_bundles()

        enabled_bundles = list(getattr(settings, "mcp_enabled_bundles", []) or [])
        for bundle_name in enabled_bundles:
            _tools_manager.enable_mcp_bundle(bundle_name)

        await _tools_manager.initialize()
        app.state.tools_manager = _tools_manager

        tools_info = _tools_manager.stats()
        logger.info(
            "engine.tools.initialized",
            total_tools=tools_info["total_tools"],
            builtin=tools_info["builtin_tools"],
            mcp=tools_info["mcp_tools"],
            mcp_bundles_initialized=tools_info["mcp_bundles_initialized"],
        )
    except Exception as e:
        logger.warning(
            "engine.tools.failed_to_initialize",
            error=str(e),
            error_type=type(e).__name__,
        )

    # === Model Discovery ===
    try:
        await start_background_refresh(interval_seconds=1800)
        logger.info("engine.model_discovery.started")
    except Exception as e:
        logger.warning("engine.model_discovery.failed_to_start", error=str(e))

    logger.info(
        "engine.started",
        uptime_ms=round((time.time() - _start_time) * 1000, 2),
    )

    yield

    # === Shutdown ===
    logger.info("engine.stopping")

    try:
        await stop_background_refresh()
        logger.info("engine.model_discovery.stopped")
    except Exception as e:
        logger.warning("engine.model_discovery.failed_to_stop", error=str(e))

    if _tools_manager is not None:
        try:
            await _tools_manager.shutdown()
            logger.info("engine.tools.stopped")
        except Exception as e:
            logger.warning("engine.tools.failed_to_stop", error=str(e))

    try:
        await shutdown_mcp_manager()
        logger.info("engine.mcp_manager.stopped")
    except Exception as e:
        logger.warning("engine.mcp_manager.failed_to_stop", error=str(e))

    try:
        await close_rag_client()
        logger.info("engine.rag_client.stopped")
    except Exception as e:
        logger.warning("engine.rag_client.failed_to_stop", error=str(e))

    try:
        await close_db()
        logger.info("engine.database.closed")
    except Exception as e:
        logger.warning("engine.database.failed_to_close", error=str(e))

    logger.info(
        "engine.stopped",
        total_uptime_seconds=round(time.time() - _start_time, 2),
    )


# ============================================================================
# App
# ============================================================================


app = FastAPI(
    title="FlowStack Engine",
    description=(
        "FlowStack Engine API — исполнительное ядро мульти-модельной LLM системы "
        "с оркестрацией агентов. "
        "Реализует MCP spec 2024-11-05: POST /api/v1/tools/mcp — unified JSON-RPC endpoint."
    ),
    version=__version__,
    lifespan=lifespan,
    docs_url="/docs" if settings.debug else None,
    redoc_url="/redoc" if settings.debug else None,
    openapi_url="/openapi.json" if settings.debug else None,
    openapi_tags=API_TAGS,
    contact={"name": "FlowStack Team", "email": "support@flowstack.dev"},
    license_info={"name": "MIT"},
)


# ============================================================================
# Middleware
# ============================================================================


@app.middleware("http")
async def add_request_id(request: Request, call_next):
    """Добавляет X-Request-ID для distributed tracing."""
    request_id = request.headers.get("X-Request-ID", str(uuid.uuid4()))
    try:
        structlog.contextvars.bind_contextvars(request_id=request_id)
    except Exception:
        pass

    response = await call_next(request)
    response.headers["X-Request-ID"] = request_id

    try:
        structlog.contextvars.clear_contextvars()
    except Exception:
        pass
    return response


app.add_middleware(AuthMiddleware, auth_mode=getattr(settings, "auth_mode", "optional"))

# ⚠️ allow_credentials=True несовместим с allow_origins=["*"] (браузер блокирует)
allowed_origins = (
    ["http://localhost:3000", "http://localhost:8080", "http://localhost:5173"]
    if settings.debug
    else [getattr(settings, "frontend_url", "http://localhost:3000")]
)

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


@app.middleware("http")
async def log_requests(request: Request, call_next):
    """Логирует все входящие запросы (только debug)."""
    if not settings.debug:
        return await call_next(request)

    start = time.time()
    logger.debug(
        "request.incoming",
        method=request.method,
        path=request.url.path,
        client=request.client.host if request.client else "unknown",
    )

    response = await call_next(request)
    duration_ms = round((time.time() - start) * 1000, 2)

    logger.debug(
        "request.completed",
        method=request.method,
        path=request.url.path,
        status_code=response.status_code,
        duration_ms=duration_ms,
    )
    return response


@app.exception_handler(Exception)
async def global_exception_handler(request: Request, exc: Exception):
    """Глобальный обработчик необработанных исключений."""
    logger.error(
        "unhandled.exception",
        method=request.method,
        path=request.url.path,
        error=str(exc),
        error_type=type(exc).__name__,
        exc_info=True,
    )
    # ⚠️ В production не раскрываем детали ошибки
    detail = str(exc) if settings.debug else "Internal server error"
    return JSONResponse(
        status_code=500,
        content={"detail": detail, "type": type(exc).__name__},
    )


# ============================================================================
# Router Registration
# ============================================================================

# Регистрация всех роутеров из src.api (chats, agents, skills, tasks, flows, mcp, tools)
register_all_routers(app)

# ✅ FIX: LLM-роутер живёт в src.llm (не в src.api) — регистрируем отдельно.
# Без этой строки ВСЕ /api/v1/llm/* возвращают 404.
app.include_router(llm_router)

logger.info(
    "engine.routers.registered",
    llm_prefix="/api/v1/llm",
)


# ============================================================================
# SCHEMAS (system endpoints)
# ============================================================================


class HealthResponse(BaseModel):
    status: str = "ok"
    version: str
    env: str
    uptime_seconds: float


class DetailedHealthResponse(BaseModel):
    status: str
    version: str
    env: str
    uptime_seconds: float
    checks: dict[str, Any]


class ReadyResponse(BaseModel):
    status: str
    models_count: int
    tools_count: int
    mcp_bundles_initialized: int


class SystemInfoResponse(BaseModel):
    version: str
    env: str
    uptime_seconds: float
    model_registry: dict[str, Any]
    tools_registry: dict[str, Any]
    api_modules: dict[str, Any]
    database: dict[str, Any]


# ============================================================================
# SYSTEM ENDPOINTS
# ============================================================================


@app.get("/health", response_model=HealthResponse, tags=["system"])
async def health() -> HealthResponse:
    """Basic health check."""
    return HealthResponse(
        status="ok",
        version=__version__,
        env=settings.app_env,
        uptime_seconds=round(time.time() - _start_time, 2) if _start_time else 0.0,
    )


@app.get("/health/detailed", response_model=DetailedHealthResponse, tags=["system"])
async def detailed_health() -> DetailedHealthResponse:
    """Deep health check — проверяет все зависимости."""
    checks: dict[str, Any] = {}
    overall_healthy = True

    # Database
    try:
        db_health = await check_database_health()
        checks["database"] = db_health
        if db_health.get("status") != "healthy":
            overall_healthy = False
    except Exception as e:
        checks["database"] = {"status": "unhealthy", "error": str(e)}
        overall_healthy = False

    # Redis
    redis_url = getattr(settings, "redis_url", None)
    if redis_url:
        try:
            import redis.asyncio as aioredis

            r = aioredis.from_url(redis_url)
            await r.ping()
            await r.aclose()
            checks["redis"] = {"status": "healthy"}
        except Exception as e:
            checks["redis"] = {"status": "unhealthy", "error": str(e)}
            overall_healthy = False

    # LLM
    try:
        model_stats = get_registry_stats()
        available_models = model_stats.get("available_models", 0)
        checks["llm"] = {
            "status": "healthy" if available_models > 0 else "degraded",
            "provider": settings.resolved_provider,
            "available_models": available_models,
        }
        if available_models == 0:
            overall_healthy = False
    except Exception as e:
        checks["llm"] = {"status": "unhealthy", "error": str(e)}
        overall_healthy = False

    # Tools
    if _tools_manager is not None:
        try:
            stats = _tools_manager.stats()
            checks["tools"] = {
                "status": "healthy",
                "total": stats["total_tools"],
                "mcp_bundles": stats["mcp_bundles_initialized"],
            }
        except Exception as e:
            checks["tools"] = {"status": "unhealthy", "error": str(e)}
            overall_healthy = False
    else:
        checks["tools"] = {"status": "degraded", "message": "Not initialized"}

    status_str = "healthy" if overall_healthy else "degraded"

    return DetailedHealthResponse(
        status=status_str,
        version=__version__,
        env=settings.app_env,
        uptime_seconds=round(time.time() - _start_time, 2) if _start_time else 0.0,
        checks=checks,
    )


@app.get("/ready", response_model=ReadyResponse, tags=["system"])
async def ready() -> ReadyResponse:
    """Readiness check. Возвращает 503 если нет доступных models."""
    try:
        model_stats = get_registry_stats()
        available_models = model_stats.get("available_models", 0)
    except Exception:
        available_models = 0

    tools_count = 0
    mcp_bundles = 0
    if _tools_manager is not None:
        stats = _tools_manager.stats()
        tools_count = stats["total_tools"]
        mcp_bundles = stats["mcp_bundles_initialized"]

    if available_models == 0:
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="No models available",
        )

    return ReadyResponse(
        status="ready",
        models_count=available_models,
        tools_count=tools_count,
        mcp_bundles_initialized=mcp_bundles,
    )


@app.get("/info", response_model=SystemInfoResponse, tags=["system"])
async def system_info() -> SystemInfoResponse:
    """System information endpoint."""
    try:
        model_stats = get_registry_stats()
    except Exception as e:
        model_stats = {"error": str(e)}

    tools_stats: dict[str, Any] = {}
    if _tools_manager is not None:
        tools_stats = _tools_manager.stats()

    api_info = get_api_info()

    db_stats: dict[str, Any] = {}
    try:
        db_stats = await get_engine_stats()
    except Exception as e:
        db_stats = {"error": str(e)}

    return SystemInfoResponse(
        version=__version__,
        env=settings.app_env,
        uptime_seconds=round(time.time() - _start_time, 2) if _start_time else 0.0,
        model_registry=model_stats,
        tools_registry=tools_stats,
        api_modules={
            "total_modules": api_info["total_modules"],
            "total_routes": api_info["total_routes"],
            "modules": [m["name"] for m in api_info["modules"]],
        },
        database=db_stats,
    )


# ============================================================================
# METRICS & CLI
# ============================================================================

_setup_prometheus(app)


def main() -> None:
    """CLI entry point."""
    logger.info(
        "cli.starting",
        host=settings.host,
        port=settings.port,
        reload=settings.debug,
        log_level=settings.log_level,
    )

    uvicorn.run(
        "src.main:app",
        host=settings.host,
        port=settings.port,
        reload=settings.debug,
        log_level=settings.log_level.lower(),
        access_log=settings.debug,
        use_colors=True,
        workers=1 if settings.debug else None,
    )


if __name__ == "__main__":
    main()