ForcePilot/backend/package/yuxi/channel/gateway/rpc_dispatcher.py

import logging
from collections.abc import AsyncGenerator, Callable, Coroutine
from typing import Any

from yuxi.channel.gateway.protocol import (
    GatewayErrorCode,
    GatewayRpcMethod,
    RpcEvent,
    RpcRequest,
    RpcResponse,
)
from yuxi.channel.gateway.rbac import GatewayRole, check_permission

logger = logging.getLogger(__name__)

RpcHandler = Callable[[RpcRequest], Coroutine[Any, Any, RpcResponse]]
RpcStreamHandler = Callable[[RpcRequest], AsyncGenerator[RpcEvent | RpcResponse, None]]


class RpcDispatcher:
    def __init__(self):
        self._handlers: dict[str, RpcHandler] = {}
        self._stream_handlers: dict[str, RpcStreamHandler] = {}
        self._lazy_handlers: dict[str, Callable[[], RpcHandler]] = {}
        self._lazy_stream_handlers: dict[str, Callable[[], RpcStreamHandler]] = {}
        self._aliases: dict[str, str] = {}
        self._enforce_rbac: bool = True

    @property
    def enforce_rbac(self) -> bool:
        return self._enforce_rbac

    @enforce_rbac.setter
    def enforce_rbac(self, value: bool) -> None:
        self._enforce_rbac = value

    def register(self, method: GatewayRpcMethod | str, handler: RpcHandler) -> None:
        key = method.value if isinstance(method, GatewayRpcMethod) else method
        if key in self._handlers or key in self._stream_handlers:
            logger.warning("RpcDispatcher: 覆盖已注册方法 %s", key)
        self._handlers[key] = handler

    def register_stream(self, method: GatewayRpcMethod | str, handler: RpcStreamHandler) -> None:
        key = method.value if isinstance(method, GatewayRpcMethod) else method
        if key in self._handlers or key in self._stream_handlers:
            logger.warning("RpcDispatcher: 覆盖已注册方法 %s", key)
        self._stream_handlers[key] = handler

    def register_lazy(self, method: GatewayRpcMethod | str, factory: Callable[[], RpcHandler]) -> None:
        key = method.value if isinstance(method, GatewayRpcMethod) else method
        if key in self._handlers or key in self._lazy_handlers:
            logger.warning("RpcDispatcher: 覆盖已有懒加载方法 %s", key)
        self._lazy_handlers[key] = factory

    def register_stream_lazy(
        self, method: GatewayRpcMethod | str, factory: Callable[[], RpcStreamHandler]
    ) -> None:
        key = method.value if isinstance(method, GatewayRpcMethod) else method
        if key in self._stream_handlers or key in self._lazy_stream_handlers:
            logger.warning("RpcDispatcher: 覆盖已有懒加载流方法 %s", key)
        self._lazy_stream_handlers[key] = factory

    def add_alias(self, alias: str, method: str) -> None:
        self._aliases[alias] = method

    def _resolve_method(self, method: str) -> str:
        return self._aliases.get(method, method)

    def _resolve_handler(self, method: str) -> RpcHandler | None:
        handler = self._handlers.get(method)
        if handler is not None:
            return handler

        factory = self._lazy_handlers.get(method)
        if factory is not None:
            handler = factory()
            self._handlers[method] = handler
            del self._lazy_handlers[method]
            return handler

        return None

    def _resolve_stream_handler(self, method: str) -> RpcStreamHandler | None:
        handler = self._stream_handlers.get(method)
        if handler is not None:
            return handler

        factory = self._lazy_stream_handlers.get(method)
        if factory is not None:
            handler = factory()
            self._stream_handlers[method] = handler
            del self._lazy_stream_handlers[method]
            return handler

        return None

    def is_stream(self, method: str) -> bool:
        resolved = self._resolve_method(method)
        return resolved in self._stream_handlers or resolved in self._lazy_stream_handlers

    async def dispatch(self, request: RpcRequest, caller_roles: list[GatewayRole] | None = None) -> RpcResponse:
        resolved_method = self._resolve_method(request.method)

        if self._enforce_rbac and caller_roles:
            primary_role = caller_roles[0] if caller_roles else None
            if not check_permission(primary_role, resolved_method):
                logger.warning(
                    "RpcDispatcher: 权限拒绝 method=%s role=%s",
                    resolved_method,
                    primary_role,
                )
                return RpcResponse(
                    id=request.id,
                    ok=False,
                    error_code=GatewayErrorCode.PERMISSION_DENIED,
                    error_message=f"方法 '{resolved_method}' 需要更高权限，当前角色为 {primary_role}",
                )

        handler = self._resolve_handler(resolved_method)
        if handler is None:
            logger.warning("RpcDispatcher: 未注册方法 %s", resolved_method)
            return RpcResponse(
                id=request.id,
                ok=False,
                error_code=GatewayErrorCode.METHOD_NOT_FOUND,
                error_message=f"方法 '{resolved_method}' 未注册",
            )
        try:
            return await handler(request)
        except Exception:
            logger.exception("RpcDispatcher: 方法 %s 执行异常", resolved_method)
            return RpcResponse(
                id=request.id,
                ok=False,
                error_code=GatewayErrorCode.INTERNAL_ERROR,
                error_message=f"方法 '{resolved_method}' 执行异常",
            )

    async def dispatch_stream(
        self, request: RpcRequest, caller_roles: list[GatewayRole] | None = None
    ) -> AsyncGenerator[RpcEvent | RpcResponse, None]:
        resolved_method = self._resolve_method(request.method)

        if self._enforce_rbac and caller_roles:
            primary_role = caller_roles[0] if caller_roles else None
            if not check_permission(primary_role, resolved_method):
                logger.warning(
                    "RpcDispatcher: 流式权限拒绝 method=%s role=%s",
                    resolved_method,
                    primary_role,
                )
                yield RpcResponse(
                    id=request.id,
                    ok=False,
                    error_code=GatewayErrorCode.PERMISSION_DENIED,
                    error_message=f"方法 '{resolved_method}' 需要更高权限，当前角色为 {primary_role}",
                )
                return

        handler = self._resolve_stream_handler(resolved_method)
        if handler is None:
            logger.warning("RpcDispatcher: 未注册流式方法 %s", resolved_method)
            yield RpcResponse(
                id=request.id,
                ok=False,
                error_code=GatewayErrorCode.METHOD_NOT_FOUND,
                error_message=f"方法 '{resolved_method}' 未注册",
            )
            return
        try:
            async for item in handler(request):
                yield item
        except Exception:
            logger.exception("RpcDispatcher: 流式方法 %s 执行异常", resolved_method)
            yield RpcResponse(
                id=request.id,
                ok=False,
                error_code=GatewayErrorCode.INTERNAL_ERROR,
                error_message=f"方法 '{resolved_method}' 执行异常",
            )

    def list_methods(self) -> list[str]:
        methods = set(self._handlers.keys())
        methods.update(self._lazy_handlers.keys())
        methods.update(self._stream_handlers.keys())
        methods.update(self._lazy_stream_handlers.keys())
        return list(methods)

    def list_aliases(self) -> dict[str, str]:
        return dict(self._aliases)


rpc_dispatcher = RpcDispatcher()