toki
a9480c5afb
- Add gateway contract for inbound queue ordering (03+02_gateway_contract) - Update receive thread model across Dart, Go, Kotlin, Python, TypeScript - Implement queue ordering logic in BaseClient and Communicator - Add comprehensive tests for queue ordering in all languages - Update documentation (PORTING_GUIDE, PROTOCOL, README) - Archive task completion logs for completed subtasks
95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import ssl
|
|
from contextlib import suppress
|
|
from typing import Any
|
|
|
|
from google.protobuf.message import DecodeError
|
|
from websockets.exceptions import ConnectionClosed
|
|
|
|
try:
|
|
from websockets.asyncio.client import connect as ws_connect
|
|
except ImportError: # websockets 12.x legacy API
|
|
from websockets import connect as ws_connect
|
|
|
|
from proto_socket.base_client import BaseClient
|
|
from proto_socket.communicator import ParserMap, type_name_of
|
|
from proto_socket.packets.message_common_pb2 import HeartBeat, PacketBase
|
|
|
|
|
|
class WsClient(BaseClient):
|
|
def __init__(
|
|
self,
|
|
ws: Any,
|
|
interval_sec: int,
|
|
wait_sec: int,
|
|
parser_map: ParserMap,
|
|
) -> None:
|
|
super().__init__(interval_sec, wait_sec, self._do_close_impl)
|
|
self._ws = ws
|
|
self._init_base()
|
|
self.communicator.initialize(self, parser_map)
|
|
self.communicator.set_write_error_handler(lambda _: asyncio.create_task(self._on_disconnected()))
|
|
self.communicator.add_listener(
|
|
type_name_of(HeartBeat),
|
|
lambda _: asyncio.create_task(self._on_heartbeat()),
|
|
)
|
|
self._read_task = asyncio.create_task(self._read_loop())
|
|
asyncio.create_task(self._send_heartbeat())
|
|
|
|
async def write_packet(self, base: PacketBase) -> None:
|
|
await self._ws.send(base.SerializeToString())
|
|
|
|
async def wait_closed(self) -> None:
|
|
await self.communicator.wait_closed()
|
|
|
|
async def _do_close_impl(self) -> None:
|
|
with suppress(Exception):
|
|
await self._ws.close()
|
|
|
|
async def _read_loop(self) -> None:
|
|
try:
|
|
async for message in self._ws:
|
|
if not isinstance(message, bytes):
|
|
continue
|
|
base = PacketBase()
|
|
base.ParseFromString(message)
|
|
await self.communicator.enqueue_inbound(
|
|
base.typeName,
|
|
base.data,
|
|
base.nonce,
|
|
base.responseNonce,
|
|
)
|
|
asyncio.create_task(self._send_heartbeat())
|
|
except (ConnectionClosed, ConnectionError, OSError, DecodeError):
|
|
pass
|
|
finally:
|
|
await self._on_disconnected()
|
|
|
|
|
|
async def connect_ws(
|
|
host: str,
|
|
port: int,
|
|
path: str,
|
|
interval_sec: int,
|
|
wait_sec: int,
|
|
parser_map: ParserMap,
|
|
) -> WsClient:
|
|
uri = f"ws://{host}:{port}{path}"
|
|
ws = await ws_connect(uri)
|
|
return WsClient(ws, interval_sec, wait_sec, parser_map)
|
|
|
|
|
|
async def connect_wss(
|
|
host: str,
|
|
port: int,
|
|
path: str,
|
|
ssl_context: ssl.SSLContext,
|
|
interval_sec: int,
|
|
wait_sec: int,
|
|
parser_map: ParserMap,
|
|
) -> WsClient:
|
|
uri = f"wss://{host}:{port}{path}"
|
|
ws = await ws_connect(uri, ssl=ssl_context)
|
|
return WsClient(ws, interval_sec, wait_sec, parser_map)
|