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.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.integrations.feishu_long_connection_core import init_feishu_long_connection_manager
from .services.scheduler_center_core import (
    AIPMOCadenceScheduler,
    KnowledgeGraphIncrementalScheduler,
    KnowledgeUrlRefreshScheduler,
    WorkloadRecalcScheduler,
)
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()


@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("数据库连接成功")

    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 = 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 = init_feishu_long_connection_manager(AsyncSessionLocal)
    await feishu_lc.start()

    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")
    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("服务关闭中")
    outbox_task.cancel()
    try:
        await outbox_task
    except asyncio.CancelledError:
        pass

    ingress_task.cancel()
    try:
        await ingress_task
    except asyncio.CancelledError:
        pass

    draft_expiry_task.cancel()
    try:
        await draft_expiry_task
    except asyncio.CancelledError:
        pass

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

    await feishu_lc.stop()
    await ai_pmo_cadence_scheduler.stop()
    await kg_scheduler.stop()
    await knowledge_scheduler.stop()
    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"})