diff --git a/gateway/relay/adapter.py b/gateway/relay/adapter.py index 5d75f956a7d6..0e3116a7ad4f 100644 --- a/gateway/relay/adapter.py +++ b/gateway/relay/adapter.py @@ -76,7 +76,25 @@ class RelayAdapter(BasePlatformAdapter): if self._transport is None: raise RuntimeError("RelayAdapter has no transport configured") self._transport.set_inbound_handler(self._on_inbound) - return await self._transport.connect() + ok = await self._transport.connect() + if not ok: + return False + # Negotiate the real capability descriptor from the connector and adopt + # it — the placeholder passed at construction is replaced by what the + # connector advertises for the platform this gateway actually fronts. + try: + descriptor = await self._transport.handshake() + except Exception as exc: # noqa: BLE001 - a failed handshake = a failed connect + logger.warning("relay handshake failed: %s", exc) + return False + self._apply_descriptor(descriptor) + return True + + def _apply_descriptor(self, descriptor: CapabilityDescriptor) -> None: + """Adopt a (re)negotiated descriptor into the live capability surface.""" + self.descriptor = descriptor + self.MAX_MESSAGE_LENGTH = descriptor.max_message_length + self.supports_code_blocks = descriptor.markdown_dialect not in ("", "plain") async def _on_inbound(self, event) -> None: """Bridge a connector-delivered MessageEvent into the normal adapter path.""" diff --git a/gateway/relay/ws_transport.py b/gateway/relay/ws_transport.py new file mode 100644 index 000000000000..75b892105ad4 --- /dev/null +++ b/gateway/relay/ws_transport.py @@ -0,0 +1,267 @@ +"""Production WebSocket RelayTransport — the gateway's live link to the connector. + +The gateway dials OUT to the connector's relay endpoint over a WebSocket and +speaks the newline-delimited JSON frame protocol defined in the connector repo +(``gateway-gateway`` ``src/relay/protocol.ts``) and mirrored in +``docs/relay-connector-contract.md``: + + gateway -> connector : hello, outbound, interrupt + connector -> gateway : descriptor, inbound, outbound_result, interrupt_inbound + +Frames: + hello {type, platform, botId} + descriptor {type, descriptor} (handshake reply) + inbound {type, event, bufferId?} (a normalized MessageEvent) + outbound {type, requestId, action} (send/edit/typing/follow_up) + outbound_result {type, requestId, result} + interrupt {type, session_key, reason?} (gateway egresses /stop) + interrupt_inbound{type, session_key, chat_id} (connector -> owning gateway) + +This is the concrete transport behind the ``RelayTransport`` Protocol; the +``RelayAdapter`` delegates all wire I/O to it. Outbound calls block on a +per-request future keyed by ``requestId`` until the matching ``outbound_result`` +arrives. A background reader task pumps inbound frames to the registered handler +and resolves pending outbound futures. + +EXPERIMENTAL: the frame schema may change without a deprecation cycle until at +least two Class-1 platforms validate it. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import uuid +from typing import Any, Dict, Optional + +from gateway.platforms.base import MessageEvent, MessageType +from gateway.session import SessionSource +from gateway.relay.descriptor import CapabilityDescriptor +from gateway.relay.transport import InboundHandler + +logger = logging.getLogger(__name__) + +try: # lazy/optional dep — mirrors gateway/platforms/feishu.py + import websockets +except ImportError: # pragma: no cover - exercised only when the extra is absent + websockets = None # type: ignore[assignment] + +WEBSOCKETS_AVAILABLE = websockets is not None + +# How long to wait for the handshake descriptor and for each outbound result. +_HANDSHAKE_TIMEOUT_S = 30.0 +_OUTBOUND_TIMEOUT_S = 30.0 + + +def _event_from_wire(raw: Dict[str, Any]) -> MessageEvent: + """Rebuild a MessageEvent from the connector's normalized inbound payload. + + The connector emits SessionSource as the snake_case wire form (§3); map it + back onto the gateway dataclasses. Unknown message types fall back to TEXT. + """ + src = raw.get("source", {}) or {} + from gateway.config import Platform + + platform = src.get("platform", "relay") + try: + platform_enum = Platform(platform) + except ValueError: + platform_enum = Platform.RELAY + + source = SessionSource( + platform=platform_enum, + chat_id=src.get("chat_id", ""), + chat_type=src.get("chat_type", "dm"), + chat_name=src.get("chat_name"), + user_id=src.get("user_id"), + user_name=src.get("user_name"), + thread_id=src.get("thread_id"), + chat_topic=src.get("chat_topic"), + user_id_alt=src.get("user_id_alt"), + chat_id_alt=src.get("chat_id_alt"), + guild_id=src.get("guild_id"), + parent_chat_id=src.get("parent_chat_id"), + message_id=src.get("message_id"), + ) + try: + msg_type = MessageType(raw.get("message_type", "text")) + except ValueError: + msg_type = MessageType.TEXT + + return MessageEvent( + text=raw.get("text", ""), + message_type=msg_type, + source=source, + message_id=raw.get("message_id"), + reply_to_message_id=raw.get("reply_to_message_id"), + media_urls=raw.get("media_urls") or [], + ) + + +class WebSocketRelayTransport: + """RelayTransport over a WebSocket connection the gateway dials to the connector.""" + + def __init__( + self, + url: str, + platform: str, + bot_id: str, + *, + connect_timeout_s: float = _HANDSHAKE_TIMEOUT_S, + outbound_timeout_s: float = _OUTBOUND_TIMEOUT_S, + ) -> None: + if not WEBSOCKETS_AVAILABLE: + raise RuntimeError( + "WebSocketRelayTransport requires the 'websockets' package " + "(install the messaging extra)." + ) + self._url = url + self._platform = platform + self._bot_id = bot_id + self._connect_timeout_s = connect_timeout_s + self._outbound_timeout_s = outbound_timeout_s + + self._ws: Any = None + self._reader: Optional[asyncio.Task[None]] = None + self._inbound: Optional[InboundHandler] = None + self._descriptor: Optional[CapabilityDescriptor] = None + self._descriptor_ready: asyncio.Future[CapabilityDescriptor] | None = None + # requestId -> future awaiting the matching outbound_result. + self._pending: Dict[str, asyncio.Future[Dict[str, Any]]] = {} + self._closing = False + + # ── lifecycle ──────────────────────────────────────────────────────── + async def connect(self) -> bool: + loop = asyncio.get_running_loop() + self._descriptor_ready = loop.create_future() + self._ws = await websockets.connect(self._url) # type: ignore[union-attr] + self._reader = asyncio.create_task(self._read_loop(), name="relay-ws-reader") + # Send hello; the descriptor arrives via the reader and resolves handshake(). + await self._send({"type": "hello", "platform": self._platform, "botId": self._bot_id}) + return True + + async def disconnect(self) -> None: + self._closing = True + if self._reader is not None: + self._reader.cancel() + try: + await self._reader + except (asyncio.CancelledError, Exception): # noqa: BLE001 - best-effort teardown + pass + self._reader = None + if self._ws is not None: + try: + await self._ws.close() + except Exception: # noqa: BLE001 + pass + self._ws = None + # Fail any in-flight outbound waiters so callers don't hang. + for fut in self._pending.values(): + if not fut.done(): + fut.set_exception(RuntimeError("relay transport closed")) + self._pending.clear() + + async def handshake(self) -> CapabilityDescriptor: + if self._descriptor is not None: + return self._descriptor + if self._descriptor_ready is None: + raise RuntimeError("handshake() called before connect()") + return await asyncio.wait_for(self._descriptor_ready, timeout=self._connect_timeout_s) + + def set_inbound_handler(self, handler: InboundHandler) -> None: + self._inbound = handler + + # ── outbound ───────────────────────────────────────────────────────── + async def send_outbound(self, action: Dict[str, Any]) -> Dict[str, Any]: + return await self._request_response(action) + + async def send_follow_up(self, action: Dict[str, Any]) -> Dict[str, Any]: + # follow_up rides the same outbound frame; the connector dispatches by + # action.op. Kept as a distinct method to satisfy the transport Protocol + # and to make the A2 call site explicit. + return await self._request_response(action) + + async def get_chat_info(self, chat_id: str) -> Dict[str, Any]: + result = await self._request_response( + {"op": "get_chat_info", "chat_id": chat_id}, frame_type="outbound" + ) + # The connector answers chat-info inside the outbound_result envelope. + info = result.get("chat_info") or result + return {"name": info.get("name", chat_id), "type": info.get("type", "dm")} + + async def send_interrupt(self, session_key: str, reason: Optional[str] = None) -> None: + await self._send({"type": "interrupt", "session_key": session_key, "reason": reason}) + + async def _request_response( + self, action: Dict[str, Any], frame_type: str = "outbound" + ) -> Dict[str, Any]: + if self._ws is None: + return {"success": False, "error": "relay transport not connected"} + request_id = uuid.uuid4().hex + loop = asyncio.get_running_loop() + fut: asyncio.Future[Dict[str, Any]] = loop.create_future() + self._pending[request_id] = fut + try: + await self._send({"type": frame_type, "requestId": request_id, "action": action}) + return await asyncio.wait_for(fut, timeout=self._outbound_timeout_s) + except asyncio.TimeoutError: + return {"success": False, "error": "relay outbound timed out"} + finally: + self._pending.pop(request_id, None) + + # ── wire I/O ───────────────────────────────────────────────────────── + async def _send(self, frame: Dict[str, Any]) -> None: + if self._ws is None: + raise RuntimeError("relay transport not connected") + await self._ws.send(json.dumps(frame) + "\n") + + async def _read_loop(self) -> None: + assert self._ws is not None + buf = "" + try: + async for chunk in self._ws: + buf += chunk if isinstance(chunk, str) else chunk.decode("utf-8") + # Newline-delimited frames; keep any trailing partial line. + *lines, buf = buf.split("\n") + for line in lines: + if line.strip(): + await self._handle_frame(line) + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 - log + let the task end; reconnection is caller policy + if not self._closing: + logger.warning("relay ws read loop ended: %s", exc) + + async def _handle_frame(self, line: str) -> None: + try: + frame = json.loads(line) + except json.JSONDecodeError: + logger.warning("relay: skipping malformed frame") + return + ftype = frame.get("type") + if ftype == "descriptor": + descriptor = CapabilityDescriptor.from_json(json.dumps(frame.get("descriptor", {}))) + self._descriptor = descriptor + if self._descriptor_ready is not None and not self._descriptor_ready.done(): + self._descriptor_ready.set_result(descriptor) + elif ftype == "inbound": + if self._inbound is not None: + event = _event_from_wire(frame.get("event", {})) + await self._inbound(event) + elif ftype == "outbound_result": + fut = self._pending.get(frame.get("requestId", "")) + if fut is not None and not fut.done(): + fut.set_result(frame.get("result", {})) + elif ftype == "interrupt_inbound": + # Bridged into the adapter's interrupt path by the runner wiring. + handler = getattr(self, "_interrupt_inbound_handler", None) + if handler is not None: + await handler(frame.get("session_key", ""), frame.get("chat_id", "")) + else: + # hello/outbound/interrupt are gateway->connector; ignore if echoed. + pass + + def set_interrupt_inbound_handler(self, handler: Any) -> None: + """Register the callback for connector->gateway interrupt_inbound frames.""" + self._interrupt_inbound_handler = handler