kaq-quant-common 0.2.12__py3-none-any.whl

kaq_quant_common/api/ws/instruction/ws_instruction_server_base.py

@@ -0,0 +1,139 @@
+import time
+from abc import ABC, abstractmethod
+
+from kaq_quant_common.api.common.api_interface import ApiInterface, api_method
+from kaq_quant_common.api.rest.instruction.models import InstructionResponseBase
+from kaq_quant_common.api.rest.instruction.models.account import (
+    ContractBalanceRequest,
+    ContractBalanceResponse,
+)
+from kaq_quant_common.api.rest.instruction.models.order import (
+    AllOpenOrdersRequest,
+    AllOpenOrdersResponse,
+    CancelOrderRequest,
+    CancelOrderResponse,
+    ChangeLeverageRequest,
+    ChangeLeverageResponse,
+    ChangeMarginModeRequest,
+    ChangeMarginModeResponse,
+    ModifyOrderRequest,
+    ModifyOrderResponse,
+    OrderRequest,
+    OrderResponse,
+    QueryMarginModeRequest,
+    QueryMarginModeResponse,
+)
+from kaq_quant_common.api.rest.instruction.models.position import (
+    QueryPositionRequest,
+    QueryPositionResponse,
+)
+from kaq_quant_common.api.rest.instruction.models.transfer import (
+    TransferRequest,
+    TransferResponse,
+)
+from kaq_quant_common.api.ws.ws_server_base import WsServerBase
+
+
+class WsInstructionServerBase(WsServerBase, ApiInterface, ABC):
+
+    def __init__(self, host: str = "0.0.0.0", port: int = 8765):
+        super().__init__(self, host, port)
+
+    # 统一处理返回数据
+    def _wrap_response(self, rsp: InstructionResponseBase):
+        if rsp is not None:
+            if rsp.event_time is None:
+                rsp.event_time = int(time.time() * 1000)
+        return rsp
+
+    # 下单
+    @api_method(OrderRequest, OrderResponse)
+    def order(self, request: OrderRequest) -> OrderResponse:
+        return self._on_order(request)
+
+    # 修改订单
+    @api_method(ModifyOrderRequest, ModifyOrderResponse)
+    def modify_order(self, request: ModifyOrderRequest) -> ModifyOrderResponse:
+        return self._on_modify_order(request)
+
+    # 撤销订单
+    @api_method(CancelOrderRequest, CancelOrderResponse)
+    def cancel_order(self, request: CancelOrderRequest) -> CancelOrderResponse:
+        return self._on_cancel_order(request)
+
+    # 查询当前全部挂单
+    @api_method(AllOpenOrdersRequest, AllOpenOrdersResponse)
+    def all_open_orders(self, request: AllOpenOrdersRequest) -> AllOpenOrdersResponse:
+        return self._on_all_open_orders(request)
+
+    # 调整杠杆
+    @api_method(ChangeLeverageRequest, ChangeLeverageResponse)
+    def change_leverage(self, request: ChangeLeverageRequest) -> ChangeLeverageResponse:
+        return self._on_change_leverage(request)
+
+    # 查询联合保证金模式
+    @api_method(QueryMarginModeRequest, QueryMarginModeResponse)
+    def query_margin_mode(self, request: QueryMarginModeRequest) -> QueryMarginModeResponse:
+        return self._on_query_margin_mode(request)
+
+    # 修改联合保证金模式
+    @api_method(ChangeMarginModeRequest, ChangeMarginModeResponse)
+    def change_margin_mode(self, request: ChangeMarginModeRequest) -> ChangeMarginModeResponse:
+        return self._on_change_margin_mode(request)
+
+    # 查询持仓
+    @api_method(QueryPositionRequest, QueryPositionResponse)
+    def query_position(self, request: QueryPositionRequest) -> QueryPositionResponse:
+        return self._on_query_position(request)
+
+    # 划转
+    @api_method(TransferRequest, TransferResponse)
+    def transfer(self, request: TransferRequest) -> TransferResponse:
+        return self._on_transfer(request)
+
+    # 查询合约账户余额
+    @api_method(ContractBalanceRequest, ContractBalanceResponse)
+    def contract_balance(self, request: ContractBalanceRequest) -> ContractBalanceResponse:
+        return self._on_contract_balance(request)
+
+    # ↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓ abstract methods
+
+    @abstractmethod
+    def _on_order(self, request: OrderRequest) -> OrderResponse:
+        pass
+
+    @abstractmethod
+    def _on_modify_order(self, request: ModifyOrderRequest) -> ModifyOrderResponse:
+        pass
+
+    @abstractmethod
+    def _on_cancel_order(self, request: CancelOrderRequest) -> CancelOrderResponse:
+        pass
+
+    @abstractmethod
+    def _on_all_open_orders(self, request: AllOpenOrdersRequest) -> AllOpenOrdersResponse:
+        pass
+
+    @abstractmethod
+    def _on_change_leverage(self, request: ChangeLeverageRequest) -> ChangeLeverageResponse:
+        pass
+
+    @abstractmethod
+    def _on_query_margin_mode(self, request: QueryMarginModeRequest) -> QueryMarginModeResponse:
+        pass
+
+    @abstractmethod
+    def _on_change_margin_mode(self, request: ChangeMarginModeRequest) -> ChangeMarginModeResponse:
+        pass
+
+    @abstractmethod
+    def _on_query_position(self, request: QueryPositionRequest) -> QueryPositionResponse:
+        pass
+
+    @abstractmethod
+    def _on_transfer(self, request: TransferRequest) -> TransferResponse:
+        pass
+
+    @abstractmethod
+    def _on_contract_balance(self, request: ContractBalanceRequest) -> ContractBalanceResponse:
+        pass

kaq_quant_common/api/ws/models.py

@@ -0,0 +1,46 @@
+from enum import Enum
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel
+
+
+# ws 消息类型
+class WsMessageType(str, Enum):
+    # 请求，对应RESPONSE
+    REQUEST = "request"
+    # 响应，对应REQUEST
+    RESPONSE = "response"
+    # 推送
+    PUSH = "push"
+    # 心跳
+    PING = "ping"
+    PONG = "pong"
+    # 订阅
+    SUBSCRIBE = "subscribe"
+    UNSUBSCRIBE = "unsubscribe"
+    #
+    ACK = "ack"
+
+
+class WsError(BaseModel):
+    code: int
+    message: str
+    details: Optional[Dict[str, Any]] = None
+
+
+class WsEnvelope(BaseModel):
+    # 请求的类型
+    type: WsMessageType
+    # 请求/响应相关
+    req_id: Optional[str] = None
+    method: Optional[str] = None
+    # 推送相关
+    topic: Optional[str] = None
+    # 负载
+    payload: Optional[Dict[str, Any]] = None
+    # 错误
+    error: Optional[WsError] = None
+
+    def model_dump_json(self) -> str:
+        # 简单包装，确保枚举序列化为字符串
+        return super().model_dump_json()

kaq_quant_common/api/ws/ws_client_base.py

@@ -0,0 +1,235 @@
+import asyncio
+import threading
+from typing import Callable, Dict, Optional, Type, TypeVar
+
+import websockets
+from pydantic import BaseModel
+
+from kaq_quant_common.api.ws.models import WsEnvelope, WsMessageType
+from kaq_quant_common.utils import logger_utils, uuid_utils
+from kaq_quant_common.api.common.auth import get_auth_token
+
+R = TypeVar("R", bound=BaseModel)
+
+
+class WsClientBase:
+    """
+    WebSocket 客户端基类：
+    - 管理长连接、自动重连和心跳（基于自定义 ping/pong）
+    - 请求/响应关联（req_id -> Future）
+    - 订阅主题并接收服务器推送
+    提供同步方法封装，便于调用。
+    """
+
+    def __init__(self, url: str, auto_reconnect: bool = True, token: Optional[str] = None):
+        self._url = url
+        self._auto_reconnect = auto_reconnect
+        self._token = token if token is not None else get_auth_token()
+        self._logger = logger_utils.get_logger(self)
+
+        # 事件循环线程
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
+        self._loop_thread: Optional[threading.Thread] = None
+
+        # 连接与控制
+        self._ws: Optional[websockets.WebSocketClientProtocol] = None
+        self._connected_event = threading.Event()
+        self._stop_flag = threading.Event()
+
+        # 请求映射
+        self._pending_requests: Dict[str, asyncio.Future] = {}
+
+        # 订阅映射：topic -> handler
+        self._subscriptions: Dict[str, Callable[[dict], None]] = {}
+
+        # 重连参数
+        self._reconnect_initial = 1.0
+        self._reconnect_max = 30.0
+
+    # ============================== 外部API ==============================
+    def connect(self):
+        """启动事件循环与连接任务"""
+        if self._loop_thread and self._loop_thread.is_alive():
+            return
+
+        # 开一条线程跑事件循环
+        def _run_loop():
+            self._loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(self._loop)
+            try:
+                self._loop.run_until_complete(self._connect_forever())
+            finally:
+                try:
+                    self._loop.close()
+                except Exception:
+                    pass
+                self._logger.warning("WsClient 事件循环结束")
+
+        # 循环接收线程，设置为守护进程，主线程关闭就自动关闭
+        self._loop_thread = threading.Thread(target=_run_loop, name="WsClientLoop", daemon=True)
+        self._loop_thread.start()
+
+    # 断开连接
+    def disconnect(self):
+        self._stop_flag.set()
+        if self._loop and self._ws:
+
+            async def _close():
+                # 关闭ws
+                try:
+                    await self._ws.close()
+                except Exception:
+                    pass
+
+            try:
+                fut = asyncio.run_coroutine_threadsafe(_close(), self._loop)
+                # 尝试快速关闭，若事件循环已结束则忽略超时
+                fut.result(timeout=1)
+            except Exception:
+                pass
+        # 等待线程结束
+        if self._loop_thread and self._loop_thread.is_alive():
+            self._loop_thread.join(timeout=5)
+
+    # 订阅，向服务器注册订阅
+    def subscribe(self, topic: str, handler: Callable[[dict], None]):
+        """注册订阅，并在连接后自动发送SUBSCRIBE"""
+        self._subscriptions[topic] = handler
+        if self._loop and self._ws:
+
+            async def _send_sub():
+                # 构造订阅消息
+                env = WsEnvelope(type=WsMessageType.SUBSCRIBE, topic=topic)
+                await self._ws.send(env.model_dump_json())
+
+            asyncio.run_coroutine_threadsafe(_send_sub(), self._loop)
+
+    # 取消订阅
+    def unsubscribe(self, topic: str):
+        self._subscriptions.pop(topic, None)
+        if self._loop and self._ws:
+
+            async def _send_unsub():
+                # 构造取消订阅消息
+                env = WsEnvelope(type=WsMessageType.UNSUBSCRIBE, topic=topic)
+                # 发送取消订阅消息
+                await self._ws.send(env.model_dump_json())
+
+            asyncio.run_coroutine_threadsafe(_send_unsub(), self._loop)
+
+    #
+    def send_request(self, method: str, request_data: BaseModel, response_model: Type[R], timeout: float = 10.0) -> R:
+        """同步请求封装，等待响应返回"""
+        if not self._connected_event.wait(timeout=5.0):
+            raise RuntimeError("WS not connected")
+
+        async def _send_and_wait() -> R:
+            # 生成请求id，用来处理响应
+            req_id = f"r_{uuid_utils.generate_uuid()}"
+            # 在事件循环中创建Future
+            fut = self._loop.create_future()
+            # 记录正在处理的请求
+            self._pending_requests[req_id] = fut
+            # 构造请求消息
+            env = WsEnvelope(type=WsMessageType.REQUEST, req_id=req_id, method=method, payload=request_data.model_dump())
+            # 发送
+            await self._ws.send(env.model_dump_json())
+            try:
+                # 等待执行完毕
+                payload: dict = await asyncio.wait_for(fut, timeout=timeout)
+                # 构造返回结构
+                return response_model(**payload)
+            finally:
+                # 移除已处理请求
+                self._pending_requests.pop(req_id, None)
+
+        # 提交到事件循环执行
+        cfut = asyncio.run_coroutine_threadsafe(_send_and_wait(), self._loop)
+        return cfut.result()
+
+    # ============================== 内部协程 ==============================
+    async def _connect_forever(self):
+        backoff = self._reconnect_initial
+        while not self._stop_flag.is_set():
+            try:
+                self._logger.info(f"WS connecting to {self._url}")
+                headers = {}
+                if self._token:
+                    headers["Authorization"] = f"Bearer {self._token}"
+                self._ws = await websockets.connect(self._url, ping_interval=None, additional_headers=headers or None)
+                self._connected_event.set()
+                self._logger.info("WS connected")
+                # 连接恢复后自动订阅
+                await self._resubscribe_all()
+                # 启动接收循环
+                await self._recv_loop()
+                # 正常退出接收循环（例如主动断开）
+                if self._stop_flag.is_set():
+                    break
+            except Exception as e:
+                self._connected_event.clear()
+                self._logger.warning(f"WS 连接失败或中断: {e}")
+                if not self._auto_reconnect or self._stop_flag.is_set():
+                    break
+                await asyncio.sleep(backoff)
+                backoff = min(backoff * 2, self._reconnect_max)
+
+        # 退出时确保关闭连接
+        try:
+            if self._ws:
+                await self._ws.close()
+        except Exception:
+            pass
+
+    # 重新自动订阅
+    async def _resubscribe_all(self):
+        for topic in list(self._subscriptions.keys()):
+            try:
+                env = WsEnvelope(type=WsMessageType.SUBSCRIBE, topic=topic)
+                await self._ws.send(env.model_dump_json())
+            except Exception as e:
+                self._logger.warning(f"自动订阅 {topic} 失败: {e}")
+
+    # 循环接收
+    async def _recv_loop(self):
+        assert self._ws is not None
+        try:
+            async for message in self._ws:
+                # 收到消息会进入到这里
+                env = None
+                try:
+                    # 校验消息格式是否正确
+                    env = WsEnvelope.model_validate_json(message)
+                except Exception as e:
+                    self._logger.error(f"解析消息失败: {e}")
+                    continue
+
+                if env.type == WsMessageType.RESPONSE:
+                    # 对应REQUEST的响应
+                    req_id = env.req_id or ""
+                    # 找到对应的请求Future
+                    fut = self._pending_requests.get(req_id)
+                    # 如果Future存在且未完成
+                    if fut and not fut.done():
+                        if env.error:
+                            # 异常响应，设置Future异常
+                            fut.set_exception(RuntimeError(env.error.get("message") if isinstance(env.error, dict) else str(env.error)))
+                        else:
+                            # 成功响应，设置Future结果
+                            fut.set_result(env.payload or {})
+                elif env.type == WsMessageType.PUSH:
+                    # 推送过来的订阅的消息
+                    topic = env.topic or ""
+                    # 找到对应订阅处理的函数
+                    handler = self._subscriptions.get(topic)
+                    if handler and env.payload is not None:
+                        # 在后台执行handler，避免阻塞事件循环
+                        asyncio.create_task(asyncio.to_thread(handler, env.payload))
+                elif env.type == WsMessageType.PING:
+                    # 心跳
+                    pong = WsEnvelope(type=WsMessageType.PONG)
+                    await self._ws.send(pong.model_dump_json())
+        except Exception as e:
+            self._logger.warning(f"WS 接收循环异常: {e}")
+        finally:
+            self._connected_event.clear()