oclaw/interfaces/ws/runtime_impl.py

"""WS runtime implementation under oclaw namespace."""

from __future__ import annotations

import asyncio
import logging
import uuid
from collections import defaultdict, deque
from typing import Any
import threading

from fastapi import WebSocket

from interfaces.gateway.dispatcher import build_gateway_method_handlers, method_names
from interfaces.ws.auth_and_hello import build_hello_ok_payload, resolve_ws_auth as resolve_ws_auth_payload
from interfaces.ws.common import (
    MAX_BUFFERED_BYTES,
    MAX_PAYLOAD_BYTES,
    PREAUTH_HANDSHAKE_TIMEOUT_MS,
    PROTOCOL_VERSION,
    TICK_INTERVAL_MS,
    WS_EVENT_REPLAY_MAX,
    WS_RATE_LIMIT_CONN_PER_WINDOW,
    WS_RATE_LIMIT_IP_PER_WINDOW,
    WS_RATE_LIMIT_USER_PER_WINDOW,
    WS_RATE_LIMIT_WINDOW_MS,
    WS_SEND_QUEUE_MAX_BYTES,
    WS_SEND_QUEUE_MAX_MESSAGES,
    error_shape as _error_shape,
    normalize_ws_attachments as _normalize_ws_attachments,
    now_ms as _now_ms,
    origin_is_allowed,
)
from interfaces.ws.events import (
    emit_agent_event as emit_agent_event_impl,
    emit_chat_event as emit_chat_event_impl,
    send_event as send_event_impl,
    send_res as send_res_impl,
)
from interfaces.ws.runtime_dispatch import dispatch_connected as dispatch_connected_impl
from interfaces.ws.runtime_helpers import handle_connect as handle_connect_impl, recv_frame as recv_frame_impl
from interfaces.ws.runtime_loop import close_ws as close_ws_impl, run_connection_loop
from interfaces.ws.server_methods_bridge import build_gateway_context, dispatch_via_server_methods
from interfaces.ws.turn_runner import run_agent_turn_via_bridge
from interfaces.ws.ws_schema import format_validation_errors, get_ws_schemas, validate_or_errors
from runtime.relay_pointer import validate_relay_share_envelope

_LOG = logging.getLogger(__name__)


class OclawWsGatewayConnection:
    _rate_lock = threading.Lock()
    _rate_by_ip: dict[str, deque[int]] = defaultdict(deque)
    _rate_by_user: dict[str, deque[int]] = defaultdict(deque)
    _stats: dict[str, int] = defaultdict(int)
    _event_buffer_by_user: dict[str, deque[dict[str, Any]]] = defaultdict(
        lambda: deque(maxlen=max(1, int(WS_EVENT_REPLAY_MAX)))
    )

    def __init__(self, ws: WebSocket):
        self.ws = ws
        self.schemas = get_ws_schemas()
        self.connected = False
        self.conn_id = uuid.uuid4().hex[:12]
        self.seq = 0
        self.started_at_ms = _now_ms()
        self.client_meta: dict[str, Any] | None = None
        self._is_webchat_client = False
        self.role = "operator"
        self.scopes: list[str] = []
        self.auth_ctx: dict[str, Any] | None = None
        self.connect_nonce = uuid.uuid4().hex
        self.handshake_failed = False
        self._subscribed_sessions_changed = False
        self._subscribed_message_keys: set[str] = set()
        self._abort_lock = threading.Lock()
        self._aborted_run_ids: set[str] = set()
        self._active_run_session: dict[str, str] = {}
        self._send_queue: asyncio.Queue[tuple[str, int]] = asyncio.Queue(maxsize=WS_SEND_QUEUE_MAX_MESSAGES)
        self._send_pending_bytes = 0
        self._send_pending_lock = asyncio.Lock()
        self._sender_task: asyncio.Task[None] | None = None
        self._event_buffer: deque[dict[str, Any]] = deque(maxlen=max(1, int(WS_EVENT_REPLAY_MAX)))
        self._rate_local: deque[int] = deque()
        self._gateway_handlers = build_gateway_method_handlers()
        self._now_ms = _now_ms
        self._error_shape = _error_shape

    async def run(self) -> None:
        self._inc_stat("ws_connections_opened")
        self._sender_task = asyncio.create_task(self._sender_loop())
        try:
            await run_connection_loop(self)
        finally:
            await self._drain_sender()
            self._inc_stat("ws_connections_closed")
            _LOG.info(
                "ws connection closed conn_id=%s total_opened=%s total_closed=%s",
                self.conn_id,
                self._stats.get("ws_connections_opened", 0),
                self._stats.get("ws_connections_closed", 0),
            )

    async def _recv_frame(self, *, preauth: bool = False) -> dict[str, Any] | None:
        return await recv_frame_impl(
            conn=self,
            preauth=preauth,
            handshake_timeout_ms=PREAUTH_HANDSHAKE_TIMEOUT_MS,
            validate_or_errors=validate_or_errors,
            format_validation_errors=format_validation_errors,
            error_shape=_error_shape,
        )

    async def _handle_connect(self, req_id: str, method: str, params: Any) -> None:
        await handle_connect_impl(
            conn=self,
            req_id=req_id,
            method=method,
            params=params,
            protocol_version=PROTOCOL_VERSION,
            validate_or_errors=validate_or_errors,
            format_validation_errors=format_validation_errors,
            error_shape=_error_shape,
        )

    async def _close_ws(self, code: int = 1000, reason: str = "done") -> None:
        await close_ws_impl(self, code=code, reason=reason)

    async def _dispatch_connected(self, req_id: str, method: str, params: Any) -> None:
        limited, bucket = self._rate_limited()
        if limited:
            self._inc_stat("ws_rate_limited")
            await self.send_res(
                req_id,
                ok=False,
                error=_error_shape("RATE_LIMITED", "too many requests", details={"bucket": bucket}),
            )
            return
        await dispatch_connected_impl(self, req_id=req_id, method=method, params=params)

    @classmethod
    def _inc_stat(cls, key: str, delta: int = 1) -> None:
        with cls._rate_lock:
            cls._stats[str(key)] = int(cls._stats.get(str(key), 0)) + int(delta)

    def mark_handshake(self, *, ok: bool) -> None:
        self._inc_stat("ws_handshake_ok" if ok else "ws_handshake_failed")
        _LOG.info(
            "ws handshake conn_id=%s ok=%s user_id=%s",
            self.conn_id,
            int(bool(ok)),
            str((self.auth_ctx or {}).get("user_id") or ""),
        )

    async def _dispatch_via_server_methods(self, *, req_id: str, method: str, params: Any) -> bool:
        return await dispatch_via_server_methods(
            req_id=req_id,
            method=method,
            params=params,
            conn_id=self.conn_id,
            is_webchat_client=bool(self._is_webchat_client),
            handlers=self._gateway_handlers,
            context=self._build_gateway_context(),
            send_res=self.send_res,
            error_shape=_error_shape,
        )

    def _build_gateway_context(self) -> dict[str, Any]:
        return build_gateway_context(
            conn_id=self.conn_id,
            subscribed_sessions_changed=bool(self._subscribed_sessions_changed),
            subscribed_message_keys=set(self._subscribed_message_keys),
            abort_lock=self._abort_lock,
            active_run_session=self._active_run_session,
            aborted_run_ids=self._aborted_run_ids,
            run_agent_turn=self.run_agent_turn,
            normalize_ws_attachments=_normalize_ws_attachments,
            validate_relay_share_envelope=validate_relay_share_envelope,
            now_ms=_now_ms,
        )

    def validate_origin(self) -> bool:
        headers = getattr(self.ws, "headers", None)
        origin = str(headers.get("origin") or "").strip() if headers is not None else ""
        host = str(headers.get("host") or "").strip() if headers is not None else ""
        allowed = origin_is_allowed(origin, host)
        if not allowed:
            _LOG.warning("ws origin blocked conn_id=%s origin=%s host=%s", self.conn_id, origin, host)
        return allowed

    def _prune_window(self, dq: deque[int], now: int) -> None:
        cutoff = now - int(WS_RATE_LIMIT_WINDOW_MS)
        while dq and dq[0] < cutoff:
            dq.popleft()

    def _rate_limited(self) -> tuple[bool, str]:
        now = _now_ms()
        self._prune_window(self._rate_local, now)
        if len(self._rate_local) >= int(WS_RATE_LIMIT_CONN_PER_WINDOW):
            return True, "connection"
        self._rate_local.append(now)

        headers = getattr(self.ws, "headers", None)
        ip = str(headers.get("x-forwarded-for") or headers.get("x-real-ip") or "").split(",")[0].strip() if headers is not None else ""
        user_id = str((self.auth_ctx or {}).get("user_id") or "").strip()
        with self._rate_lock:
            if ip:
                ip_bucket = self._rate_by_ip[ip]
                self._prune_window(ip_bucket, now)
                if len(ip_bucket) >= int(WS_RATE_LIMIT_IP_PER_WINDOW):
                    return True, "ip"
                ip_bucket.append(now)
            if user_id:
                user_bucket = self._rate_by_user[user_id]
                self._prune_window(user_bucket, now)
                if len(user_bucket) >= int(WS_RATE_LIMIT_USER_PER_WINDOW):
                    return True, "user"
                user_bucket.append(now)
        return False, ""

    async def _queue_send_text(self, text: str) -> bool:
        payload = str(text or "")
        payload_size = len(payload.encode("utf-8", errors="ignore"))
        async with self._send_pending_lock:
            if self._send_pending_bytes + payload_size > int(WS_SEND_QUEUE_MAX_BYTES):
                _LOG.warning("ws send queue bytes exceeded conn_id=%s", self.conn_id)
                return False
            self._send_pending_bytes += payload_size
        try:
            self._send_queue.put_nowait((payload, payload_size))
            return True
        except asyncio.QueueFull:
            async with self._send_pending_lock:
                self._send_pending_bytes = max(0, self._send_pending_bytes - payload_size)
            _LOG.warning("ws send queue full conn_id=%s", self.conn_id)
            return False

    async def _sender_loop(self) -> None:
        while True:
            item = await self._send_queue.get()
            if item[0] == "__STOP__":
                self._send_queue.task_done()
                return
            payload, payload_size = item
            try:
                await self.ws.send_text(payload)
            except Exception:
                return
            finally:
                async with self._send_pending_lock:
                    self._send_pending_bytes = max(0, self._send_pending_bytes - payload_size)
                self._send_queue.task_done()

    async def _drain_sender(self) -> None:
        if self._sender_task is None:
            return
        try:
            self._send_queue.put_nowait(("__STOP__", 0))
        except Exception:
            pass
        try:
            await self._sender_task
        except Exception:
            pass
        self._sender_task = None

    def remember_event(self, frame: dict[str, Any]) -> None:
        snap = dict(frame or {})
        self._event_buffer.append(snap)
        user_id = str((self.auth_ctx or {}).get("user_id") or "").strip()
        if not user_id:
            return
        with self._rate_lock:
            bucket = self._event_buffer_by_user.get(user_id)
            if bucket is None or bucket.maxlen != max(1, int(WS_EVENT_REPLAY_MAX)):
                bucket = deque(maxlen=max(1, int(WS_EVENT_REPLAY_MAX)))
                self._event_buffer_by_user[user_id] = bucket
            bucket.append(dict(snap))

    async def replay_events_since(self, seq: int) -> None:
        after = int(seq or 0)
        frames = list(self._event_buffer)
        user_id = str((self.auth_ctx or {}).get("user_id") or "").strip()
        if user_id:
            with self._rate_lock:
                shared = list(self._event_buffer_by_user.get(user_id) or [])
            if shared:
                frames = shared
        for frame in frames:
            fseq = int(frame.get("seq") or 0)
            if fseq <= after:
                continue
            text = frame.get("_raw")
            if not isinstance(text, str) or not text:
                continue
            await self._queue_send_text(text)

    def build_hello_ok(self, _connect_params: dict[str, Any] | None) -> dict[str, Any]:
        methods = ["connect", *method_names()]
        return build_hello_ok_payload(
            conn_id=self.conn_id,
            started_at_ms=int(self.started_at_ms or _now_ms()),
            role=str(self.role or "operator"),
            scopes=list(self.scopes or []),
            protocol_version=PROTOCOL_VERSION,
            methods=methods,
            max_payload_bytes=MAX_PAYLOAD_BYTES,
            max_buffered_bytes=MAX_BUFFERED_BYTES,
            tick_interval_ms=TICK_INTERVAL_MS,
            now_ms=_now_ms(),
        )

    async def send_res(self, req_id: str, *, ok: bool, payload: Any | None = None, error: Any | None = None) -> None:
        await send_res_impl(self, req_id, ok=ok, payload=payload, error=error)

    async def send_event(self, event: str, payload: Any | None = None) -> None:
        await send_event_impl(self, event, payload)

    async def emit_agent_event(self, *, run_id: str, stream: str, data: dict[str, Any]) -> None:
        await emit_agent_event_impl(self, run_id=run_id, stream=stream, data=data, now_ms=_now_ms())

    async def emit_chat_event(
        self,
        *,
        run_id: str,
        state: str,
        delta: str = "",
        reply: str = "",
        error: str = "",
        message: dict[str, Any] | None = None,
        session_key: str | None = None,
        seq: int | None = None,
    ) -> None:
        await emit_chat_event_impl(
            self,
            run_id=run_id,
            state=state,
            delta=delta,
            reply=reply,
            error=error,
            message=message,
            session_key=session_key,
            seq=seq,
        )

    def resolve_ws_auth(self, connect_params: dict[str, Any] | None) -> dict[str, Any]:
        return resolve_ws_auth_payload(connect_params)

    async def run_agent_turn(self, req_id: str, p: dict[str, Any], *, session_id: str, send_response: bool = True) -> None:
        await run_agent_turn_via_bridge(
            conn=self,
            req_id=req_id,
            p=p,
            session_id=session_id,
            send_response=send_response,
            normalize_ws_attachments=_normalize_ws_attachments,
            validate_relay_share_envelope=validate_relay_share_envelope,
            now_ms=_now_ms,
            error_shape=_error_shape,
        )


async def ws_gateway_loop(ws: WebSocket) -> None:
    await OclawWsGatewayConnection(ws).run()


__all__ = ["OclawWsGatewayConnection", "ws_gateway_loop"]