ForcePilot/backend/package/yuxi/channel/bootstrap.py

"""多渠道网关启动与关闭封装。

将 `lifespan.py` 中分散的渠道组件初始化、组装与销毁逻辑集中到此处，
降低应用启动器与渠道子系统内部实现的耦合。
"""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from yuxi.channel.config import ChannelConfigManager
from yuxi.channel.lifecycle import ChannelLifecycleManager
from yuxi.channel.message.dedupe import MessageDeduper
from yuxi.channel.message.dispatcher import ChannelGateway
from yuxi.channel.middlewares.inbound import (
    CreateRunMiddleware,
    DedupeMiddleware,
    RouteMiddleware,
    SecurityMiddleware,
    SessionMiddleware,
    TransactionMiddleware,
)
from yuxi.channel.middlewares.outbound import (
    BuildMessageMiddleware,
    ChunkMiddleware,
    DowngradeMiddleware,
    EnrichMiddleware,
    FormatMiddleware,
    MediaUploadMiddleware,
    SendMiddleware,
    StatusUpdateMiddleware,
)
from yuxi.channel.middlewares.registry import (
    InboundMiddlewareRegistry,
    OutboundMiddlewareRegistry,
)
from yuxi.channel.outbound import OutboundDispatcher
from yuxi.channel.outbound.retry import compensate_channel_messages
from yuxi.channel.plugins.loader import load_plugins
from yuxi.channel.plugins.registry import get_registry
from yuxi.channel.routing.router import BindingRouter
from yuxi.channel.security.allowlist import AllowlistChecker
from yuxi.channel.security.bot_loop import BotLoopDetector
from yuxi.channel.security.identity import IdentityLinkResolver
from yuxi.channel.security.pairing import PairingChecker, PairingManager
from yuxi.channel.security.policy import SecurityPolicy
from yuxi.channel.security.rate_limit import RateLimiter
from yuxi.channel.security.registry import SecurityCheckerRegistry
from yuxi.channel.session.manager import SessionManager
from yuxi.services.run_queue_service import get_redis_client
from yuxi.utils.logging_config import logger

if TYPE_CHECKING:
    from fastapi import FastAPI


class ChannelGatewayBootstrap:
    """多渠道网关启动器：负责加载插件、组装网关、启停生命周期与后台补偿任务。"""

    def __init__(self) -> None:
        self._lifecycle_manager: ChannelLifecycleManager | None = None
        self._outbound_dispatcher: OutboundDispatcher | None = None
        self._gateway: ChannelGateway | None = None
        self._security_checker_registry: SecurityCheckerRegistry | None = None
        self._inbound_registry: InboundMiddlewareRegistry | None = None
        self._outbound_registry: OutboundMiddlewareRegistry | None = None
        self._compensate_task: asyncio.Task[None] | None = None
        self._started: bool = False

    async def start(self, app: FastAPI) -> None:
        """加载插件、组装并启动多渠道网关，将关键组件挂载到 ``app.state``。

        启动过程中任意步骤失败都会触发回滚（``stop``），避免组件泄漏。
        重复调用会被忽略。
        """
        if self._started:
            logger.warning("Channel gateway already started, ignoring duplicate start")
            return

        load_plugins()
        registry = get_registry()
        config_manager = ChannelConfigManager()

        outbound_registry = OutboundMiddlewareRegistry()
        self._outbound_registry = outbound_registry

        dispatcher = OutboundDispatcher(
            registry=registry,
            config_manager=config_manager,
            outbound_registry=outbound_registry,
        )
        self._outbound_dispatcher = dispatcher

        outbound_registry.register(BuildMessageMiddleware())
        outbound_registry.register(MediaUploadMiddleware())
        outbound_registry.register(DowngradeMiddleware())
        outbound_registry.register(ChunkMiddleware())
        outbound_registry.register(FormatMiddleware())
        outbound_registry.register(EnrichMiddleware())
        outbound_registry.register(SendMiddleware())
        outbound_registry.register(
            StatusUpdateMiddleware(update_status=dispatcher._finalize_dispatch)
        )

        inbound_registry = InboundMiddlewareRegistry()
        self._inbound_registry = inbound_registry

        router = BindingRouter()
        session_manager = SessionManager(router)
        redis = await get_redis_client()

        identity_resolver = await IdentityLinkResolver.create()

        security_checker_registry = SecurityCheckerRegistry()
        security_checker_registry.register(AllowlistChecker())
        security_checker_registry.register(PairingChecker(PairingManager()))
        security_checker_registry.register(BotLoopDetector())
        security_checker_registry.register(RateLimiter(redis))
        self._security_checker_registry = security_checker_registry

        security_policy = SecurityPolicy(security_checker_registry, identity_resolver)

        lifecycle_manager = ChannelLifecycleManager(registry, config_manager)
        self._lifecycle_manager = lifecycle_manager

        dedupe = MessageDeduper()

        gateway = ChannelGateway(
            registry=registry,
            config_manager=config_manager,
            session_manager=session_manager,
            binding_router=router,
            outbound_dispatcher=dispatcher,
            security_policy=security_policy,
            lifecycle_manager=lifecycle_manager,
            inbound_registry=inbound_registry,
            dedupe=dedupe,
        )
        self._gateway = gateway

        inbound_registry.register(DedupeMiddleware(dedupe))
        inbound_registry.register(SecurityMiddleware(security_policy, dedupe))
        inbound_registry.register(TransactionMiddleware(dedupe))
        inbound_registry.register(SessionMiddleware(session_manager))
        inbound_registry.register(RouteMiddleware(router))
        inbound_registry.register(
            CreateRunMiddleware(
                create_run=gateway.create_agent_run,
                record_message=session_manager.record_message,
            )
        )

        # 必须在 lifecycle 启动前注册 handler，否则 WebSocket/Polling 类渠道
        # 在 start_all() 与 set_message_handler() 之间收到的消息会丢失。
        lifecycle_manager.set_message_handler(lambda raw, ct, aid: gateway.on_transport_message(raw, ct, aid))

        try:
            await security_checker_registry.start_all()
            await inbound_registry.start_all()
            await lifecycle_manager.start_all()
            await dispatcher.start()
            await gateway.start_config_change_listener()

            app.state.channel_lifecycle_manager = lifecycle_manager
            app.state.channel_outbound_dispatcher = dispatcher
            app.state.channel_gateway = gateway

            self._compensate_task = asyncio.create_task(self._compensate_loop())
            self._started = True
        except Exception:
            logger.exception("Failed to start channel gateway, rolling back")
            await self.stop()
            raise

    async def stop(self) -> None:
        """按依赖顺序关闭多渠道网关相关组件与后台任务。

        重复调用或 ``start`` 未成功时调用均为安全空操作；
        若 ``start`` 中途失败，也会按实际已创建的组件进行清理。
        """
        if (
            not self._started
            and self._lifecycle_manager is None
            and self._gateway is None
            and self._outbound_dispatcher is None
            and self._security_checker_registry is None
        ):
            return

        self._started = False

        if self._security_checker_registry is not None:
            try:
                await self._security_checker_registry.stop_all()
            except Exception:
                logger.exception("Failed to stop security checker registry")
            finally:
                self._security_checker_registry = None

        if self._inbound_registry is not None:
            try:
                await self._inbound_registry.stop_all()
            except Exception:
                logger.exception("Failed to stop inbound middleware registry")
            finally:
                self._inbound_registry = None

        if self._compensate_task is not None:
            try:
                self._compensate_task.cancel()
                try:
                    await self._compensate_task
                except asyncio.CancelledError:
                    pass
            except Exception:
                logger.exception("Failed to stop channel compensate task")
            finally:
                self._compensate_task = None

        # gateway.stop() 内部已会关闭 lifecycle，避免重复停止。
        if self._gateway is not None:
            try:
                await self._gateway.stop()
            except Exception:
                logger.exception("Failed to stop channel gateway")
        elif self._lifecycle_manager is not None:
            try:
                await self._lifecycle_manager.stop_all()
            except Exception:
                logger.exception("Failed to stop channel lifecycle manager")

        if self._outbound_dispatcher is not None:
            try:
                await self._outbound_dispatcher.stop()
            except Exception:
                logger.exception("Failed to stop channel outbound dispatcher")

        self._gateway = None
        self._lifecycle_manager = None
        self._outbound_dispatcher = None

    async def _compensate_loop(self) -> None:
        while True:
            await asyncio.sleep(60)
            try:
                await compensate_channel_messages()
            except Exception:
                logger.exception("Channel compensate loop error")