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

import json as _json
import time as _time
import uuid as _uuid
from dataclasses import dataclass, field
from enum import StrEnum


class FrameType(StrEnum):
    REQUEST = "request"
    RESPONSE = "response"
    EVENT = "event"


class GatewayErrorCode(StrEnum):
    INVALID_REQUEST = "INVALID_REQUEST"
    METHOD_NOT_FOUND = "METHOD_NOT_FOUND"
    INVALID_PARAMS = "INVALID_PARAMS"
    INTERNAL_ERROR = "INTERNAL_ERROR"
    AUTH_ERROR = "AUTH_ERROR"
    PERMISSION_DENIED = "PERMISSION_DENIED"
    RATE_LIMITED = "RATE_LIMITED"
    TIMEOUT = "TIMEOUT"
    NOT_CONNECTED = "NOT_CONNECTED"
    UNAVAILABLE = "UNAVAILABLE"
    DEVICE_IDENTITY_REQUIRED = "DEVICE_IDENTITY_REQUIRED"


class GatewayRpcMethod(StrEnum):
    CONNECT = "connect"
    CHAT_SEND = "chat.send"
    CHAT_STREAM = "chat.stream"
    CHANNELS_LIST = "channels.list"
    CHANNELS_STATUS = "channels.status"
    CHANNELS_CONFIGURE = "channels.configure"
    SESSIONS_LIST = "sessions.list"
    SESSIONS_HISTORY = "sessions.history"
    PLUGINS_LIST = "plugins.list"
    PLUGINS_INSTALL = "plugins.install"
    SYSTEM_HEALTH = "system.health"
    SYSTEM_CONFIG = "system.config"
    SYSTEM_VERSION = "system.version"
    CHANNELS_START = "channels.start"
    CHANNELS_STOP = "channels.stop"
    CHANNELS_RESTART = "channels.restart"
    PAIRING_LIST = "pairing.list"
    PAIRING_APPROVE = "pairing.approve"
    PAIRING_REJECT = "pairing.reject"
    ALLOWLIST_GET = "allowlist.get"
    ALLOWLIST_ADD = "allowlist.add"
    ALLOWLIST_REMOVE = "allowlist.remove"
    CHAT_HISTORY = "chat.history"
    CHAT_CANCEL = "chat.cancel"
    CRON_LIST = "cron.list"
    CRON_CREATE = "cron.create"
    CRON_UPDATE = "cron.update"
    CRON_DELETE = "cron.delete"
    CRON_FORCE_RUN = "cron.force_run"
    CRON_PAUSE = "cron.pause"
    CRON_RESUME = "cron.resume"
    CRON_DIAGNOSTICS = "cron.diagnostics"
    CRON_RUN_LOG = "cron.run_log"
    AGENT_TOOLS_LIST = "agentTools.list"
    AGENT_TOOLS_EXECUTE = "agentTools.execute"
    SYSTEM_LOG_TAIL = "system.log.tail"
    SESSIONS_CREATE = "sessions.create"
    SESSIONS_DELETE = "sessions.delete"
    CONFIG_GET = "config.get"
    CONFIG_SET = "config.set"
    MESSAGE_SEND = "message.send"
    MESSAGE_ACTION = "message.action"
    CHANNELS_PROBE = "channels.probe"
    IDENTITY_LINKS_LIST = "identity_links.list"
    IDENTITY_LINKS_ADD = "identity_links.add"
    IDENTITY_LINKS_REMOVE = "identity_links.remove"


class DeliveryMode(StrEnum):
    DIRECT = "direct"
    GATEWAY = "gateway"
    HYBRID = "hybrid"


@dataclass
class RpcFrame:
    frame_type: FrameType


@dataclass
class RpcRequest(RpcFrame):
    frame_type: FrameType = FrameType.REQUEST
    id: str = field(default_factory=lambda: _uuid.uuid4().hex[:12])
    method: str = ""
    params: dict | None = None
    session_id: str | None = None
    caller_user_id: str | None = None


@dataclass
class RpcResponse(RpcFrame):
    frame_type: FrameType = FrameType.RESPONSE
    id: str = ""
    ok: bool = True
    result: dict | None = None
    error_code: GatewayErrorCode | None = None
    error_message: str | None = None
    error_data: dict | None = None
    retryable: bool = False
    retry_after_ms: int | None = None


@dataclass
class HelloOk(RpcResponse):
    frame_type: FrameType = FrameType.RESPONSE
    id: str = ""
    ok: bool = True
    result: dict | None = field(default_factory=lambda: {
        "protocolVersion": 1,
        "server": {"version": "1.0.0", "connId": ""},
        "features": {"methods": [], "events": []},
        "auth": {"role": "", "scopes": []},
        "policy": {
            "maxPayload": 25 * 1024 * 1024,
            "maxBufferedBytes": 256 * 1024,
            "tickIntervalMs": 30_000,
        },
    })


@dataclass
class RpcEvent(RpcFrame):
    frame_type: FrameType = FrameType.EVENT
    event: str = ""
    data: dict | None = None
    timestamp: float = field(default_factory=_time.time)
    seq: int | None = None
    state_version: int | None = None


def marshal_frame(frame: RpcFrame) -> str:
    if isinstance(frame, RpcRequest):
        payload = {"type": frame.frame_type.value, "id": frame.id, "method": frame.method}
        if frame.params is not None:
            payload["params"] = frame.params
        if frame.session_id is not None:
            payload["sessionId"] = frame.session_id
    elif isinstance(frame, RpcResponse):
        payload = {"type": frame.frame_type.value, "id": frame.id, "ok": frame.ok}
        if frame.result is not None:
            payload["result"] = frame.result
        if frame.error_code is not None:
            error_payload = {
                "code": frame.error_code.value,
                "message": frame.error_message or "",
            }
            if frame.error_data is not None:
                error_payload["data"] = frame.error_data
            if frame.retryable:
                error_payload["retryable"] = True
            if frame.retry_after_ms is not None:
                error_payload["retryAfterMs"] = frame.retry_after_ms
            payload["error"] = error_payload
    elif isinstance(frame, RpcEvent):
        payload = {"type": frame.frame_type.value, "event": frame.event, "timestamp": frame.timestamp}
        if frame.data is not None:
            payload["data"] = frame.data
        if frame.seq is not None:
            payload["seq"] = frame.seq
        if frame.state_version is not None:
            payload["stateVersion"] = frame.state_version
    else:
        raise TypeError(f"Unknown frame type: {type(frame)}")
    return _json.dumps(payload, ensure_ascii=False, default=str)


def unmarshal_frame(data: str | bytes) -> RpcFrame:
    if isinstance(data, bytes):
        data = data.decode("utf-8")
    raw = _json.loads(data)
    frame_type = raw.get("type", FrameType.REQUEST.value)
    if frame_type == FrameType.REQUEST.value:
        return RpcRequest(
            id=raw.get("id", ""),
            method=raw.get("method", ""),
            params=raw.get("params"),
            session_id=raw.get("sessionId"),
        )
    elif frame_type == FrameType.RESPONSE.value:
        error = raw.get("error")
        return RpcResponse(
            id=raw.get("id", ""),
            ok=raw.get("ok", True),
            result=raw.get("result"),
            error_code=GatewayErrorCode(error["code"]) if error else None,
            error_message=error.get("message") if error else None,
            error_data=error.get("data") if error else None,
            retryable=error.get("retryable", False) if error else False,
            retry_after_ms=error.get("retryAfterMs") if error else None,
        )
    elif frame_type == FrameType.EVENT.value:
        return RpcEvent(
            event=raw.get("event", ""),
            data=raw.get("data"),
            timestamp=raw.get("timestamp", _time.time()),
            seq=raw.get("seq"),
            state_version=raw.get("stateVersion"),
        )
    raise ValueError(f"Unknown frame type: {frame_type}")