ForcePilot/backend/server/utils/lifespan.py

import os
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver

from yuxi.services.task_service import tasker
from yuxi.agents.mcp.service import ensure_builtin_mcp_servers_in_db
from yuxi.models.providers.service import ensure_builtin_model_providers_in_db
from yuxi.services.run_queue_service import close_queue_clients, get_redis_client
from yuxi.storage.postgres.manager import pg_manager
from yuxi.external_systems import register_adapters, shutdown_external_systems
from yuxi.knowledge import knowledge_base
from yuxi.utils import logger
from yuxi.agents.backends.sandbox import init_sandbox_provider, shutdown_sandbox_provider
from yuxi import get_version
from server.utils.inflight_tracker import InflightRequestTracker


@asynccontextmanager
async def lifespan(app: FastAPI):

    # 初始化在途请求计数器（InflightTrackingMiddleware 依赖）
    app.state.inflight_tracker = InflightRequestTracker()

    # 注册外部系统适配器（避免模块导入副作用）
    try:
        register_adapters()
    except Exception as e:
        logger.error(f"Failed to register external system adapters during startup: {e}")

    """FastAPI lifespan事件管理器"""
    # 初始化数据库连接
    try:
        pg_manager.initialize()
        await pg_manager.create_tables()
        await pg_manager.ensure_business_schema()
        await pg_manager.ensure_knowledge_schema()
        await pg_manager.ensure_external_schema()
        await pg_manager.ensure_scheduler_schema()
    except Exception as e:
        logger.error(f"Failed to initialize database during startup: {e}")

    # =========================================================
    # 多渠道网关（channels）启动编排
    # 装配 DI 容器并执行 HostBootstrap.bootstrap()，将容器挂载至
    # app.state.channel_di_container 供路由层 Depends 解析。
    # 失败记录错误但不终止应用启动：channels 路由会返回 InternalError，
    # 其他模块照常工作。关停时若已装配则执行 HostShutdown.shutdown()。
    # =========================================================
    host_shutdown = None
    try:
        from yuxi.channels.adapters import _logger_adapter as channel_logger
        from yuxi.channels.infrastructure import (
            DependencyInjectionContainer,
            create_host_bootstrap,
            create_host_shutdown,
        )

        channel_di_container = DependencyInjectionContainer()
        plugin_dir = str(
            Path(__file__).resolve().parents[2]
            / "package"
            / "yuxi"
            / "channels"
            / "plugins"
        )
        host_bootstrap = await create_host_bootstrap(
            ensure_schema=pg_manager,
            logger=channel_logger,
            plugin_dir=plugin_dir,
            di_container=channel_di_container,
        )
        # 提前创建 host_shutdown，bootstrap 失败时也需执行关停释放资源
        host_shutdown = create_host_shutdown(host_bootstrap, channel_di_container)
        try:
            await host_bootstrap.bootstrap()
            app.state.channel_di_container = channel_di_container
            logger.info("Channels module bootstrap complete")
        except Exception as e:
            logger.error(f"Channels module bootstrap failed: {e}")
            try:
                await host_shutdown.shutdown()
            except Exception as shutdown_err:
                logger.error(f"Channels module shutdown after bootstrap failure failed: {shutdown_err}")
            host_shutdown = None
    except Exception as e:
        logger.error(f"Failed to initialize channels module during startup: {e}")

    # 确保内置 MCP 服务器定义存在于数据库
    try:
        await ensure_builtin_mcp_servers_in_db()
    except Exception as e:
        logger.error(f"Failed to ensure builtin MCP servers during startup: {e}")

    try:
        from yuxi.agents.skills.service import init_builtin_skills

        async with pg_manager.get_async_session_context() as session:
            await init_builtin_skills(session)
    except Exception as e:
        logger.error(f"Failed to initialize builtin skills during startup: {e}")

    try:
        from yuxi.repositories.agent_repository import AgentRepository

        async with pg_manager.get_async_session_context() as session:
            repository = AgentRepository(session)
            await repository.ensure_default_agent()
            await repository.ensure_general_purpose_subagent()
            await repository.ensure_web_search_subagent()
            await repository.ensure_deep_research_agents()
    except Exception as e:
        logger.error(f"Failed to ensure default agent during startup: {e}")

    # 初始化内置模型供应商配置
    try:
        async with pg_manager.get_async_session_context() as session:
            await ensure_builtin_model_providers_in_db(session)
    except Exception as e:
        logger.error(f"Failed to ensure builtin model providers during startup: {e}")

    # 初始化模型缓存（v2 模型选择使用）
    try:
        from yuxi.models.providers.cache import model_cache
        from yuxi.models.providers.service import get_all_model_providers

        async with pg_manager.get_async_session_context() as session:
            providers = await get_all_model_providers(session)
            model_cache.rebuild(providers)
    except Exception as e:
        logger.error(f"Failed to initialize model cache during startup: {e}")

    # 初始化知识库管理器
    if os.environ.get("LITE_MODE", "").lower() in ("true", "1"):
        logger.info("LITE_MODE enabled, skipping knowledge base initialization")
    else:
        try:
            await knowledge_base.initialize()
        except Exception as e:
            logger.error(f"Failed to initialize knowledge base manager: {e}")

    # 预热 Redis（run 队列）
    try:
        redis = await get_redis_client()
        await redis.ping()
    except Exception as e:
        logger.warning(f"Run queue redis unavailable on startup: {e}")

    try:
        init_sandbox_provider()
    except Exception as e:
        logger.error(f"Failed to initialize sandbox provider during startup: {e}")

    # =========================================================
    # 2. 核心修复：在这里执行一次 setup()，建完表就拉倒
    # =========================================================
    checkpointer = AsyncPostgresSaver(pg_manager.langgraph_pool)
    await checkpointer.setup()
    print("LangGraph Checkpoint tables verified/created!")

    await tasker.start()
    logger.info(f"""

░██     ░██                       ░██
 ░██   ░██
  ░██ ░██   ░██    ░██ ░██    ░██ ░██
   ░████    ░██    ░██  ░██  ░██  ░██
    ░██     ░██    ░██   ░█████   ░██
    ░██     ░██   ░███  ░██  ░██  ░██
    ░██      ░█████░██ ░██    ░██ ░██  v{get_version()}

    """)
    logger.info("Yuxi backend startup complete")
    yield
    # 关停多渠道网关编排层（仅当 bootstrap 成功装配时）
    if host_shutdown is not None:
        try:
            await host_shutdown.shutdown()
        except Exception as e:
            logger.error(f"Channels module shutdown failed: {e}")
    await tasker.shutdown()
    await shutdown_external_systems()
    shutdown_sandbox_provider()
    await close_queue_clients()
    await pg_manager.close()