OpenOPC/opc/layer4_tools/collaboration_rpc.py

"""Local RPC transport for ``opc-collab`` calls.

The external agent still invokes the normal ``opc-collab`` CLI. When OpenOPC
spawns that agent, the broker exposes a short-lived local endpoint and injects
its address/token into the environment. The CLI then sends the collaboration
tool call to the already-running broker, so database writes stay in the host
runtime instead of inside the agent sandbox.
"""

from __future__ import annotations

import asyncio
import contextlib
import errno
import json
import os
import select
import secrets
import shutil
import socket
import tempfile
import uuid
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal


OPC_COLLAB_RPC_PATH = "OPC_COLLAB_RPC_PATH"
OPC_COLLAB_RPC_TOKEN = "OPC_COLLAB_RPC_TOKEN"
OPC_COLLAB_RPC_TRANSPORT = "OPC_COLLAB_RPC_TRANSPORT"
OPC_COLLAB_RPC_HOST = "OPC_COLLAB_RPC_HOST"
OPC_COLLAB_RPC_PORT = "OPC_COLLAB_RPC_PORT"
_RPC_TIMEOUT_SECONDS = 30.0
_RPC_MAX_BYTES = 16 * 1024 * 1024

DispatchCallable = Callable[[str, dict[str, Any]], Awaitable[tuple[dict[str, Any], bool]]]
RpcTransport = Literal["auto", "fifo", "tcp"]


def _infrastructure_error(message: str, *, tool_name: str = "") -> dict[str, Any]:
    payload = {
        "error": str(message or "collaboration broker RPC failed"),
        "error_type": "infrastructure",
        "retryable": True,
    }
    if tool_name:
        payload["tool_name"] = tool_name
    return payload


def fifo_rpc_supported() -> bool:
    """Return whether this runtime can create POSIX FIFOs."""
    return os.name != "nt" and callable(getattr(os, "mkfifo", None))


def default_collaboration_rpc_transport() -> Literal["fifo", "tcp"]:
    return "fifo" if fifo_rpc_supported() else "tcp"


def resolve_collaboration_rpc_transport(
    transport: str | None = "auto",
) -> Literal["fifo", "tcp"]:
    normalized = str(transport or "auto").strip().lower()
    if normalized in {"", "auto"}:
        return default_collaboration_rpc_transport()
    if normalized == "fifo":
        if not fifo_rpc_supported():
            raise RuntimeError("FIFO collaboration RPC is unavailable on this platform")
        return "fifo"
    if normalized == "tcp":
        return "tcp"
    raise ValueError(f"Unsupported collaboration RPC transport: {transport}")


def rpc_env_available(env: Mapping[str, str] | None = None) -> bool:
    source = env if env is not None else os.environ
    token = str(source.get(OPC_COLLAB_RPC_TOKEN, "")).strip()
    if not token:
        return False
    transport = str(source.get(OPC_COLLAB_RPC_TRANSPORT, "")).strip().lower()
    if not transport:
        # Legacy FIFO environment from older brokers.
        return bool(str(source.get(OPC_COLLAB_RPC_PATH, "")).strip())
    if transport == "fifo":
        return bool(str(source.get(OPC_COLLAB_RPC_PATH, "")).strip())
    if transport == "tcp":
        host = str(source.get(OPC_COLLAB_RPC_HOST, "")).strip()
        raw_port = str(source.get(OPC_COLLAB_RPC_PORT, "")).strip()
        try:
            port = int(raw_port)
        except ValueError:
            return False
        return bool(host) and 0 < port <= 65535
    return False


def rpc_env_configured(env: Mapping[str, str] | None = None) -> bool:
    source = env if env is not None else os.environ
    return any(
        str(source.get(key, "")).strip()
        for key in (
            OPC_COLLAB_RPC_TRANSPORT,
            OPC_COLLAB_RPC_PATH,
            OPC_COLLAB_RPC_HOST,
            OPC_COLLAB_RPC_PORT,
            OPC_COLLAB_RPC_TOKEN,
        )
    )


def _json_line(payload: dict[str, Any]) -> bytes:
    return json.dumps(payload, ensure_ascii=False, default=str).encode("utf-8") + b"\n"


def _decode_rpc_response(raw: bytes, *, tool_name: str) -> tuple[dict[str, Any], bool]:
    try:
        response = json.loads(raw.decode("utf-8"))
    except json.JSONDecodeError as exc:
        return _infrastructure_error(f"collaboration broker RPC returned invalid JSON: {exc}", tool_name=tool_name), True
    if not isinstance(response, dict):
        return _infrastructure_error("collaboration broker RPC returned a non-object response", tool_name=tool_name), True

    result = response.get("result")
    normalized = result if isinstance(result, dict) else {"result": result}
    is_error = bool(response.get("is_error")) or "error" in normalized
    return normalized, is_error


def _request_payload(
    tool_name: str,
    args: dict[str, Any],
    *,
    token: str,
    response_path: str = "",
) -> dict[str, Any]:
    request = {
        "token": token,
        "tool_name": str(tool_name or "").strip(),
        "args": dict(args or {}),
    }
    if response_path:
        request["response_path"] = response_path
    return request


def _write_fifo_nonblocking(path: Path, payload: dict[str, Any]) -> None:
    data = _json_line(payload)
    if len(data) > _RPC_MAX_BYTES:
        raise ValueError("collaboration RPC request exceeds max payload size")
    try:
        fd = os.open(path, os.O_WRONLY | os.O_NONBLOCK)
    except OSError as exc:
        if exc.errno == errno.ENXIO:
            raise RuntimeError("collaboration broker RPC is not accepting requests") from exc
        raise
    try:
        view = memoryview(data)
        while view:
            try:
                written = os.write(fd, view)
                view = view[written:]
            except BlockingIOError:
                _readable, writable, _errors = select.select([], [fd], [], _RPC_TIMEOUT_SECONDS)
                if not writable:
                    raise TimeoutError("collaboration broker RPC write timed out")
    finally:
        os.close(fd)


async def _read_fifo_response(fd: int, *, timeout_seconds: float) -> bytes:
    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout_seconds
    chunks = bytearray()
    while loop.time() < deadline:
        try:
            chunk = os.read(fd, 65536)
        except BlockingIOError:
            await asyncio.sleep(0.02)
            continue
        if chunk:
            chunks.extend(chunk)
            if b"\n" in chunk:
                line, _sep, _rest = bytes(chunks).partition(b"\n")
                return line + b"\n"
            if len(chunks) > _RPC_MAX_BYTES:
                raise RuntimeError("collaboration broker RPC response exceeds max payload size")
        else:
            await asyncio.sleep(0.02)
    raise TimeoutError("collaboration broker RPC response timed out")


async def call_collaboration_rpc(
    tool_name: str,
    args: dict[str, Any],
    *,
    env: Mapping[str, str] | None = None,
) -> tuple[dict[str, Any], bool]:
    """Call the broker-owned collaboration RPC endpoint from ``opc-collab``."""
    source = env if env is not None else os.environ
    transport = str(source.get(OPC_COLLAB_RPC_TRANSPORT, "")).strip().lower() or "fifo"
    if transport == "fifo":
        return await _call_fifo_collaboration_rpc(tool_name, args, env=source)
    if transport == "tcp":
        return await _call_tcp_collaboration_rpc(tool_name, args, env=source)
    return _infrastructure_error(
        f"collaboration broker RPC transport is unsupported: {transport}",
        tool_name=tool_name,
    ), True


async def _call_fifo_collaboration_rpc(
    tool_name: str,
    args: dict[str, Any],
    *,
    env: Mapping[str, str],
) -> tuple[dict[str, Any], bool]:
    token = str(env.get(OPC_COLLAB_RPC_TOKEN, "")).strip()
    raw_request_path = str(env.get(OPC_COLLAB_RPC_PATH, "")).strip()
    if not raw_request_path or not token:
        return _infrastructure_error("collaboration broker RPC is not configured", tool_name=tool_name), True
    if not fifo_rpc_supported():
        return _infrastructure_error(
            "FIFO collaboration broker RPC is unavailable on this platform",
            tool_name=tool_name,
        ), True
    request_path = Path(raw_request_path)

    response_dir = request_path.parent / "responses"
    response_path = response_dir / f"{uuid.uuid4().hex}.fifo"
    response_fd: int | None = None
    try:
        response_dir.mkdir(parents=True, exist_ok=True)
        os.mkfifo(response_path, 0o600)
        response_fd = os.open(response_path, os.O_RDONLY | os.O_NONBLOCK)
        request = _request_payload(
            tool_name,
            args,
            token=token,
            response_path=str(response_path),
        )
        _write_fifo_nonblocking(request_path, request)
        raw = await _read_fifo_response(response_fd, timeout_seconds=_RPC_TIMEOUT_SECONDS)
    except Exception as exc:
        return _infrastructure_error(f"collaboration broker RPC failed: {exc}", tool_name=tool_name), True
    finally:
        if response_fd is not None:
            with contextlib.suppress(OSError):
                os.close(response_fd)
        with contextlib.suppress(FileNotFoundError):
            response_path.unlink()

    return _decode_rpc_response(raw, tool_name=tool_name)


async def _call_tcp_collaboration_rpc(
    tool_name: str,
    args: dict[str, Any],
    *,
    env: Mapping[str, str],
) -> tuple[dict[str, Any], bool]:
    host = str(env.get(OPC_COLLAB_RPC_HOST, "")).strip()
    raw_port = str(env.get(OPC_COLLAB_RPC_PORT, "")).strip()
    token = str(env.get(OPC_COLLAB_RPC_TOKEN, "")).strip()
    if not host or not raw_port or not token:
        return _infrastructure_error("collaboration broker RPC is not configured", tool_name=tool_name), True
    try:
        port = int(raw_port)
    except ValueError:
        return _infrastructure_error(
            f"collaboration broker RPC port is invalid: {raw_port}",
            tool_name=tool_name,
        ), True
    if not 0 < port <= 65535:
        return _infrastructure_error(
            f"collaboration broker RPC port is out of range: {port}",
            tool_name=tool_name,
        ), True

    request = _request_payload(tool_name, args, token=token)
    data = _json_line(request)
    if len(data) > _RPC_MAX_BYTES:
        return _infrastructure_error(
            "collaboration RPC request exceeds max payload size",
            tool_name=tool_name,
        ), True

    writer: asyncio.StreamWriter | None = None
    try:
        reader, writer = await asyncio.wait_for(
            asyncio.open_connection(host=host, port=port, limit=_RPC_MAX_BYTES + 1024),
            timeout=_RPC_TIMEOUT_SECONDS,
        )
        writer.write(data)
        await asyncio.wait_for(writer.drain(), timeout=_RPC_TIMEOUT_SECONDS)
        raw = await _read_stream_line(reader, timeout_seconds=_RPC_TIMEOUT_SECONDS)
    except Exception as exc:
        return _infrastructure_error(f"collaboration broker RPC failed: {exc}", tool_name=tool_name), True
    finally:
        if writer is not None:
            writer.close()
            with contextlib.suppress(Exception):
                await writer.wait_closed()

    return _decode_rpc_response(raw, tool_name=tool_name)


@dataclass
class CollaborationRpcServer:
    transport: Literal["fifo", "tcp"]
    token: str
    request_path: Path | None = None
    rpc_dir: Path | None = None
    request_fd: int | None = None
    task: asyncio.Task[None] | None = None
    tcp_server: asyncio.AbstractServer | None = None
    host: str = ""
    port: int = 0

    @property
    def client_env(self) -> dict[str, str]:
        if self.transport == "tcp":
            return {
                OPC_COLLAB_RPC_TRANSPORT: "tcp",
                OPC_COLLAB_RPC_HOST: self.host,
                OPC_COLLAB_RPC_PORT: str(self.port),
                OPC_COLLAB_RPC_TOKEN: self.token,
            }
        return {
            OPC_COLLAB_RPC_PATH: str(self.request_path or ""),
            OPC_COLLAB_RPC_TOKEN: self.token,
            OPC_COLLAB_RPC_TRANSPORT: "fifo",
        }

    async def close(self) -> None:
        if self.tcp_server is not None:
            self.tcp_server.close()
            with contextlib.suppress(Exception):
                await self.tcp_server.wait_closed()
        if self.task is not None:
            self.task.cancel()
            with contextlib.suppress(asyncio.CancelledError):
                await self.task
        if self.request_fd is not None:
            with contextlib.suppress(OSError):
                os.close(self.request_fd)
        if self.rpc_dir is not None:
            shutil.rmtree(self.rpc_dir, ignore_errors=True)


async def start_collaboration_rpc_server(
    dispatch: DispatchCallable,
    *,
    transport_parent: str | os.PathLike[str] | None = None,
    transport: RpcTransport = "auto",
) -> CollaborationRpcServer | None:
    """Start a broker-local collaboration RPC server."""
    resolved_transport = resolve_collaboration_rpc_transport(transport)
    if resolved_transport == "tcp":
        return await _start_tcp_collaboration_rpc_server(dispatch)
    return await _start_fifo_collaboration_rpc_server(dispatch, transport_parent=transport_parent)


async def _start_fifo_collaboration_rpc_server(
    dispatch: DispatchCallable,
    *,
    transport_parent: str | os.PathLike[str] | None = None,
) -> CollaborationRpcServer:
    parent = Path(transport_parent) if transport_parent else Path(tempfile.gettempdir())
    rpc_dir = Path(tempfile.mkdtemp(prefix="openopc-collab-rpc-", dir=str(parent)))
    request_path = rpc_dir / "requests.fifo"
    token = secrets.token_urlsafe(32)
    os.mkfifo(request_path, 0o600)
    request_fd = os.open(request_path, os.O_RDONLY | os.O_NONBLOCK)

    async def _serve() -> None:
        buffer = bytearray()
        while True:
            try:
                chunk = os.read(request_fd, 65536)
            except BlockingIOError:
                await asyncio.sleep(0.02)
                continue
            if not chunk:
                await asyncio.sleep(0.02)
                continue
            buffer.extend(chunk)
            if len(buffer) > _RPC_MAX_BYTES:
                buffer.clear()
                continue
            while True:
                newline_index = buffer.find(b"\n")
                if newline_index < 0:
                    break
                raw = bytes(buffer[:newline_index])
                del buffer[: newline_index + 1]
                await _handle_request_line(raw)

    async def _handle_request_line(raw: bytes) -> None:
        if not raw:
            return
        try:
            request = json.loads(raw.decode("utf-8"))
        except json.JSONDecodeError:
            return
        if not isinstance(request, dict):
            return
        response = await _handle_rpc_request(request, dispatch=dispatch, token=token)
        await _respond(request, response)

    task = asyncio.create_task(_serve())
    return CollaborationRpcServer(
        transport="fifo",
        request_path=request_path,
        rpc_dir=rpc_dir,
        token=token,
        request_fd=request_fd,
        task=task,
    )


async def _start_tcp_collaboration_rpc_server(dispatch: DispatchCallable) -> CollaborationRpcServer:
    token = secrets.token_urlsafe(32)

    async def _handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        try:
            raw = await _read_stream_line(reader, timeout_seconds=_RPC_TIMEOUT_SECONDS)
            try:
                request = json.loads(raw.decode("utf-8"))
            except json.JSONDecodeError as exc:
                response = {
                    "result": _infrastructure_error(
                        f"collaboration broker RPC received invalid JSON: {exc}",
                    ),
                    "is_error": True,
                }
            else:
                if isinstance(request, dict):
                    response = await _handle_rpc_request(request, dispatch=dispatch, token=token)
                else:
                    response = {
                        "result": _infrastructure_error("collaboration broker RPC received a non-object request"),
                        "is_error": True,
                    }
            data = _json_line(response)
            if len(data) > _RPC_MAX_BYTES:
                data = _json_line(
                    {
                        "result": _infrastructure_error("collaboration broker RPC response exceeds max payload size"),
                        "is_error": True,
                    }
                )
            writer.write(data)
            await writer.drain()
        except Exception:
            with contextlib.suppress(Exception):
                writer.write(
                    _json_line(
                        {
                            "result": _infrastructure_error("collaboration broker RPC connection failed"),
                            "is_error": True,
                        }
                    )
                )
                await writer.drain()
        finally:
            writer.close()
            with contextlib.suppress(Exception):
                await writer.wait_closed()

    server = await asyncio.start_server(
        _handle_client,
        host="127.0.0.1",
        port=0,
        family=socket.AF_INET,
        limit=_RPC_MAX_BYTES + 1024,
    )
    sockets = server.sockets or []
    if not sockets:
        server.close()
        await server.wait_closed()
        raise RuntimeError("collaboration RPC TCP server did not expose a listening socket")
    host, port = sockets[0].getsockname()[:2]
    return CollaborationRpcServer(
        transport="tcp",
        token=token,
        tcp_server=server,
        host=str(host),
        port=int(port),
    )


async def _read_stream_line(
    reader: asyncio.StreamReader,
    *,
    timeout_seconds: float,
) -> bytes:
    buffer = bytearray()
    while True:
        chunk = await asyncio.wait_for(reader.read(65536), timeout=timeout_seconds)
        if not chunk:
            if buffer:
                return bytes(buffer)
            raise RuntimeError("collaboration broker RPC connection closed before a response")
        buffer.extend(chunk)
        if len(buffer) > _RPC_MAX_BYTES:
            raise RuntimeError("collaboration broker RPC payload exceeds max size")
        newline_index = buffer.find(b"\n")
        if newline_index >= 0:
            return bytes(buffer[: newline_index + 1])


async def _handle_rpc_request(
    request: dict[str, Any],
    *,
    dispatch: DispatchCallable,
    token: str,
) -> dict[str, Any]:
    tool_name = str(request.get("tool_name", "") or "").strip()
    if str(request.get("token", "")) != token:
        return {
            "result": _infrastructure_error("collaboration RPC token rejected", tool_name=tool_name),
            "is_error": True,
        }
    raw_args = request.get("args")
    tool_args = raw_args if isinstance(raw_args, dict) else {}
    try:
        result, is_error = await dispatch(tool_name, tool_args)
        return {"result": result, "is_error": bool(is_error)}
    except Exception as exc:
        return {
            "result": _infrastructure_error(
                f"collaboration broker RPC failed: {exc}",
                tool_name=tool_name,
            ),
            "is_error": True,
        }


async def _respond(request: dict[str, Any], response: dict[str, Any]) -> None:
    response_path = Path(str(request.get("response_path", "") or "").strip())
    if not response_path:
        return
    with contextlib.suppress(Exception):
        _write_fifo_nonblocking(response_path, response)