PMFlow/backend/app/main.py

"""PMFlow FastAPI 应用入口。"""

from __future__ import annotations

import asyncio
import os
from contextlib import asynccontextmanager, suppress

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from .api.router import api_router
from .core.config import settings
from .core.logging import configure_logging, get_logger
from .core.response import success_response
from .core.database import AsyncSessionLocal, engine
from .core.startup.db_guard import check_database_dialect_guard
from .core.startup.ai_pmo_ingress_async import ai_pmo_ingress_drain_loop
from .core.startup.context_memory_worker import context_memory_compaction_drain_loop, stop_context_memory_worker
from .core.startup.draft_expiry_worker import draft_expiry_drain_loop
from .core.startup.feishu_outbox import feishu_outbox_drain_loop
from .core.startup.model_probe import probe_model_connectivity_on_startup, startup_log
from .core.startup.world_model_warmup import warmup_world_models
from .middleware import AuditContextMiddleware
from .services.ai_orchestration_run import AIOrchestrationRunService
from .services.context_cache_service import register_context_cache_invalidation_hooks
from .services.skill import init_system_skill_tags
from .services.semantic_fact_service import register_semantic_fact_hooks

configure_logging(settings.debug)
logger = get_logger(__name__)
register_context_cache_invalidation_hooks()
register_semantic_fact_hooks()


def _startup_step_enabled(step: str, enabled: bool, env_name: str) -> bool:
    if enabled:
        return True
    startup_log(f"跳过{step}", reason=f"{env_name}=false")
    return False


@asynccontextmanager
async def lifespan(_: FastAPI):
    """应用生命周期：初始化依赖并启动后台任务。"""
    backend_port = str(os.getenv("BACKEND_PORT", "8000") or "8000").strip() or "8000"
    startup_log("服务启动初始化中")

    auto_create_tables = os.getenv("AUTO_CREATE_TABLES", "false").lower() in {"1", "true", "yes"}
    if auto_create_tables:
        logger.error(
            "runtime_ddl_blocked",
            auto_create_tables=auto_create_tables,
            reason="runtime ddl is forbidden, run alembic migration instead",
        )
        raise RuntimeError("AUTO_CREATE_TABLES=true 已被禁用；请使用 Alembic migration 变更数据库。")

    async with engine.begin() as conn:
        await conn.run_sync(check_database_dialect_guard)
    startup_log("数据库连接成功")

    if _startup_step_enabled("模型连通性探测", settings.startup_model_probe_enabled, "STARTUP_MODEL_PROBE_ENABLED"):
        await probe_model_connectivity_on_startup()

    async with AsyncSessionLocal() as db:
        try:
            created = await init_system_skill_tags(db, tenant_id=1)
            if created:
                logger.info("system_skill_tags_seeded", created=created)
        except Exception as exc:  # pragma: no cover
            logger.warning("seed_system_skills_failed", error=str(exc))

        try:
            interrupted = await AIOrchestrationRunService(db).mark_stale_runs_interrupted(older_than_minutes=10)
            if interrupted:
                await db.commit()
                logger.info("stale_ai_orchestration_runs_interrupted", count=interrupted)
        except Exception as exc:  # pragma: no cover
            logger.warning("interrupt_stale_ai_orchestration_runs_failed", error=str(exc))

    scheduler = None
    knowledge_scheduler = None
    kg_scheduler = None
    ai_pmo_cadence_scheduler = None
    if _startup_step_enabled("定时调度器", settings.schedulers_enabled, "STARTUP_SCHEDULERS_ENABLED"):
        from .services.scheduler_center_core import (
            AIPMOCadenceScheduler,
            KnowledgeGraphIncrementalScheduler,
            KnowledgeUrlRefreshScheduler,
            WorkloadRecalcScheduler,
        )

        scheduler = WorkloadRecalcScheduler(AsyncSessionLocal)
        scheduler.start()
        knowledge_scheduler = KnowledgeUrlRefreshScheduler(AsyncSessionLocal)
        knowledge_scheduler.start()
        kg_scheduler = KnowledgeGraphIncrementalScheduler(AsyncSessionLocal)
        kg_scheduler.start()
        ai_pmo_cadence_scheduler = AIPMOCadenceScheduler(AsyncSessionLocal)
        ai_pmo_cadence_scheduler.start()

    feishu_lc = None
    if _startup_step_enabled("飞书长连接", settings.feishu_long_connection_enabled, "FEISHU_LONG_CONNECTION_ENABLED"):
        from .services.integrations.feishu_long_connection_core import init_feishu_long_connection_manager

        feishu_lc = init_feishu_long_connection_manager(AsyncSessionLocal)
        await feishu_lc.start()

    outbox_task = None
    ingress_task = None
    draft_expiry_task = None
    context_memory_task = None
    if _startup_step_enabled("后台队列消费", settings.background_workers_enabled, "STARTUP_BACKGROUND_WORKERS_ENABLED"):
        outbox_task = asyncio.create_task(feishu_outbox_drain_loop(), name="feishu-outbox-drain-loop")
        ingress_task = asyncio.create_task(ai_pmo_ingress_drain_loop(), name="ai-pmo-ingress-drain-loop")
        draft_expiry_task = asyncio.create_task(draft_expiry_drain_loop(), name="draft-expiry-drain-loop")
        context_memory_task = asyncio.create_task(
            context_memory_compaction_drain_loop(tenant_id=1),
            name="context-memory-compaction-drain-loop",
        )

    wm_warmup_task = None
    if _startup_step_enabled("World Model 预热", settings.world_model_warmup_enabled, "STARTUP_WORLD_MODEL_WARMUP_ENABLED"):
        wm_warmup_task = asyncio.create_task(warmup_world_models(), name="wm-world-model-warmup")
    startup_log("后端运行中", url=f"http://127.0.0.1:{backend_port}")

    yield

    startup_log("服务关闭中")
    if outbox_task is not None:
        outbox_task.cancel()
        try:
            await outbox_task
        except asyncio.CancelledError:
            pass

    if ingress_task is not None:
        ingress_task.cancel()
        try:
            await ingress_task
        except asyncio.CancelledError:
            pass

    if draft_expiry_task is not None:
        draft_expiry_task.cancel()
        try:
            await draft_expiry_task
        except asyncio.CancelledError:
            pass

    await stop_context_memory_worker(context_memory_task)

    if wm_warmup_task is not None and not wm_warmup_task.done():
        wm_warmup_task.cancel()
        with suppress(asyncio.CancelledError):
            await wm_warmup_task

    if feishu_lc is not None:
        await feishu_lc.stop()
    if ai_pmo_cadence_scheduler is not None:
        await ai_pmo_cadence_scheduler.stop()
    if kg_scheduler is not None:
        await kg_scheduler.stop()
    if knowledge_scheduler is not None:
        await knowledge_scheduler.stop()
    if scheduler is not None:
        await scheduler.stop()
    await engine.dispose()
    startup_log("服务已关闭")


app = FastAPI(
    title=settings.app_name,
    debug=settings.debug,
    version="0.1.0",
    lifespan=lifespan,
)

app.add_middleware(AuditContextMiddleware)

origins = [item.strip() for item in settings.cors_origins.split(",") if item.strip()]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins or ["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(api_router)


@app.get("/healthz")
async def healthz() -> dict:
    return success_response({"status": "ok"})