ForcePilot/backend/package/yuxi/channel/extensions/signal/client.py

import asyncio
import base64
import json
import uuid
from collections.abc import AsyncGenerator
from dataclasses import dataclass
from typing import Any

import aiohttp

MAX_SINGLE_RESPONSE_BYTES = 1_048_576
MAX_SSE_BUFFER_BYTES = 1_048_576
MAX_SSE_EVENT_DATA_BYTES = 1_048_576


class SignalRpcError(Exception):
    pass


@dataclass
class SignalSseEvent:
    event: str
    data: str


class SignalRpcClient:
    def __init__(self, base_url: str, timeout: float = 10.0):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout

    async def call(self, method: str, params: dict | None = None, account: str | None = None) -> Any:
        body = {
            "jsonrpc": "2.0",
            "method": method,
            "params": params or {},
            "id": str(uuid.uuid4()),
        }

        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout)) as session:
            async with session.post(
                f"{self.base_url}/api/v1/rpc",
                json=body,
                headers={"Content-Type": "application/json"},
            ) as resp:
                if resp.status == 201:
                    return None

                raw = await resp.content.read()
                if not raw:
                    raise SignalRpcError("Signal RPC empty response")

                try:
                    data = json.loads(raw)
                except Exception:
                    raise SignalRpcError("Signal RPC returned malformed JSON")

                if "error" in data and data["error"]:
                    err = data["error"]
                    raise SignalRpcError(f"Signal RPC {err.get('code', -1)}: {err.get('message', 'unknown')}")

                if "result" not in data:
                    raise SignalRpcError("Signal RPC returned invalid response envelope")

                return data["result"]

    async def check(self, timeout_ms: int = 1000) -> bool:
        try:
            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=timeout_ms / 1000)) as session:
                async with session.get(f"{self.base_url}/api/v1/check") as resp:
                    return resp.status == 200
        except Exception:
            return False

    async def version(self) -> str:
        result = await self.call("version")
        if isinstance(result, dict):
            return str(result.get("version", "unknown"))
        return str(result)

    async def get_attachment(
        self,
        attachment_id: str,
        *,
        account: str | None = None,
        recipient: str | None = None,
        group_id: str | None = None,
    ) -> bytes:
        params: dict = {"id": attachment_id}
        if account:
            params["account"] = account
        if recipient:
            params["recipient"] = recipient
        if group_id:
            params["groupId"] = group_id

        result = await self.call("getAttachment", params, account=account)

        if isinstance(result, str):
            return base64.b64decode(result)
        if isinstance(result, dict) and "data" in result:
            return base64.b64decode(result["data"])
        raise SignalRpcError("getAttachment returned unexpected format")


class SignalSseClient:
    def __init__(self, base_url: str, account: str | None = None):
        self.base_url = base_url.rstrip("/")
        self.account = account

    async def stream_events(self, cancel: asyncio.Event | None = None) -> AsyncGenerator[SignalSseEvent, None]:
        url = f"{self.base_url}/api/v1/events"
        params = {}
        if self.account:
            params["account"] = self.account

        async with aiohttp.ClientSession() as session:
            async with session.get(url, params=params) as resp:
                event_type = None
                data_lines: list[str] = []
                total_bytes = 0

                async for line in resp.content:
                    if cancel and cancel.is_set():
                        break

                    decoded = line.decode("utf-8").rstrip("\r\n")

                    if decoded == "":
                        if data_lines:
                            yield SignalSseEvent(
                                event=event_type or "message",
                                data="\n".join(data_lines),
                            )
                        event_type = None
                        data_lines.clear()
                        total_bytes = 0
                    elif decoded.startswith(":"):
                        continue
                    elif decoded.startswith("event:"):
                        event_type = decoded[6:].strip()
                    elif decoded.startswith("data:"):
                        data_str = decoded[5:].strip()
                        data_lines.append(data_str)
                        total_bytes += len(data_str)
                        if total_bytes > MAX_SSE_BUFFER_BYTES:
                            raise RuntimeError("SSE buffer size exceeded")
                    elif decoded.startswith("id:"):
                        pass