auto_trade_sys/trading_system/book_ticker_stream.py

"""
最优挂单 WebSocket 流：订阅 !bookTicker，维护全市场最优买/卖价缓存。
用于滑点估算、入场价格优化，提升交易执行效果。
文档：更新速度 5s，推送所有交易对的最优挂单（最高买单、最低卖单）。
"""
import asyncio
import json
import logging
import time
from typing import Dict, Optional, Any

logger = logging.getLogger(__name__)

# 最优挂单缓存：symbol -> { bidPrice, bidQty, askPrice, askQty, time }
_book_ticker_cache: Dict[str, Dict[str, Any]] = {}
_book_ticker_updated_at: float = 0.0


def get_book_ticker_cache() -> Dict[str, Dict[str, Any]]:
    """返回当前最优挂单缓存。"""
    return dict(_book_ticker_cache)


def get_book_ticker(symbol: str) -> Optional[Dict[str, Any]]:
    """获取指定交易对的最优挂单；无缓存时返回 None。"""
    return _book_ticker_cache.get(symbol.upper())


def is_book_ticker_cache_fresh(max_age_sec: float = 10.0) -> bool:
    """缓存是否在 max_age_sec 秒内更新过且非空。"""
    if not _book_ticker_cache:
        return False
    return (time.monotonic() - _book_ticker_updated_at) <= max_age_sec


def estimate_slippage(symbol: str, side: str, quantity: float) -> Optional[float]:
    """
    估算滑点（USDT）：基于最优挂单估算执行价格与标记价格的偏差。
    
    Args:
        symbol: 交易对
        side: BUY/SELL
        quantity: 下单数量
    
    Returns:
        估算滑点（USDT），None 表示无法估算
    """
    ticker = get_book_ticker(symbol)
    if not ticker:
        return None
    try:
        bid_price = float(ticker.get("bidPrice", 0))
        ask_price = float(ticker.get("askPrice", 0))
        bid_qty = float(ticker.get("bidQty", 0))
        ask_qty = float(ticker.get("askQty", 0))
        if bid_price <= 0 or ask_price <= 0:
            return None
        mid_price = (bid_price + ask_price) / 2
        if side == "BUY":
            # 买单：用 askPrice（卖一）估算，滑点 = (askPrice - midPrice) * quantity
            if ask_qty >= quantity:
                slippage = (ask_price - mid_price) * quantity
            else:
                # 数量超过卖一，需要吃多档，粗略估算
                slippage = (ask_price - mid_price) * quantity * 1.2
        else:  # SELL
            # 卖单：用 bidPrice（买一）估算，滑点 = (midPrice - bidPrice) * quantity
            if bid_qty >= quantity:
                slippage = (mid_price - bid_price) * quantity
            else:
                slippage = (mid_price - bid_price) * quantity * 1.2
        return slippage
    except Exception:
        return None


class BookTickerStream:
    """订阅合约 !bookTicker，持续更新 _book_ticker_cache。无需 listenKey，公开行情。"""

    def __init__(self, testnet: bool = False):
        self.testnet = testnet
        self._ws = None
        self._task: Optional[asyncio.Task] = None
        self._running = False

    def _ws_url(self) -> str:
        if self.testnet:
            return "wss://stream.binancefuture.com/ws/!bookTicker"
        return "wss://fstream.binance.com/ws/!bookTicker"

    async def start(self) -> bool:
        if self._running:
            return True
        self._running = True
        self._task = asyncio.create_task(self._run_ws())
        logger.info("BookTickerStream: 已启动（!bookTicker），用于滑点估算")
        return True

    async def stop(self):
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
            self._ws = None
        logger.info("BookTickerStream: 已停止")

    async def _run_ws(self):
        import aiohttp
        while self._running:
            url = self._ws_url()
            try:
                async with aiohttp.ClientSession() as session:
                    async with session.ws_connect(
                        url, heartbeat=50, timeout=aiohttp.ClientTimeout(total=15)
                    ) as ws:
                        self._ws = ws
                        logger.info("BookTickerStream: WS 已连接")
                        async for msg in ws:
                            if not self._running:
                                break
                            if msg.type == aiohttp.WSMsgType.TEXT:
                                self._handle_message(msg.data)
                            elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSE):
                                break
            except asyncio.CancelledError:
                break
            except Exception as e:
                err_msg = getattr(e, "message", str(e)) or repr(e)
                err_type = type(e).__name__
                logger.warning(
                    "BookTickerStream: WS 异常 %s: %s，10s 后重连",
                    err_type, err_msg,
                    exc_info=logger.isEnabledFor(logging.DEBUG),
                )
                await asyncio.sleep(10)
            self._ws = None
            if not self._running:
                break

    def _handle_message(self, raw: str):
        global _book_ticker_cache, _book_ticker_updated_at
        try:
            data = json.loads(raw)
        except Exception:
            return
        # 可能是单条对象或组合流格式
        if isinstance(data, dict) and "stream" in data:
            ticker_data = data.get("data", {})
        else:
            ticker_data = data

        if not isinstance(ticker_data, dict) or ticker_data.get("e") != "bookTicker":
            return

        s = (ticker_data.get("s") or "").strip().upper()
        if not s or not s.endswith("USDT"):
            return

        try:
            _book_ticker_cache[s] = {
                "symbol": s,
                "bidPrice": float(ticker_data.get("b", 0)),
                "bidQty": float(ticker_data.get("B", 0)),
                "askPrice": float(ticker_data.get("a", 0)),
                "askQty": float(ticker_data.get("A", 0)),
                "time": int(ticker_data.get("T", 0)),
            }
        except (TypeError, ValueError):
            return

        _book_ticker_updated_at = time.monotonic()
        logger.debug(f"BookTickerStream: 已更新 {s} bid={_book_ticker_cache[s]['bidPrice']:.4f} ask={_book_ticker_cache[s]['askPrice']:.4f}")