230 lines
8.1 KiB
Python
230 lines
8.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
User WebSocket Client
|
|
|
|
Subscribes to Polymarket User Channel for order/trade tracking.
|
|
Used to verify order execution before retry.
|
|
|
|
Docs: https://docs.polymarket.com/developers/CLOB/websocket/user-channel
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import websockets
|
|
from typing import Optional, Dict, Callable, Any
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
|
|
from src.proxy_util import apply_proxy_env, ws_connect, ws_connect_kwargs as _ws_connect_kwargs
|
|
|
|
apply_proxy_env()
|
|
|
|
logger = logging.getLogger("btc_live.user_ws")
|
|
|
|
WS_URL = "wss://ws-subscriptions-clob.polymarket.com/ws/user"
|
|
|
|
|
|
@dataclass
|
|
class OrderStatus:
|
|
"""Tracks order status from WebSocket."""
|
|
order_id: str
|
|
asset_id: str
|
|
side: str
|
|
price: float
|
|
original_size: int
|
|
size_matched: int = 0
|
|
status: str = "PENDING" # PENDING, PLACED, MATCHED, CANCELLED
|
|
trades: list = field(default_factory=list)
|
|
timestamp: datetime = field(default_factory=datetime.now)
|
|
|
|
|
|
class UserWebSocket:
|
|
"""
|
|
WebSocket client for User Channel.
|
|
|
|
Tracks order placements and fills in real-time.
|
|
"""
|
|
|
|
def __init__(self, api_key: str, api_secret: str = "", api_passphrase: str = ""):
|
|
self.api_key = api_key
|
|
self.api_secret = api_secret
|
|
self.api_passphrase = api_passphrase
|
|
self._ws = None
|
|
self._ws_cm = None
|
|
self._connected = False
|
|
self._running = False
|
|
self._orders: Dict[str, OrderStatus] = {}
|
|
|
|
# Callbacks
|
|
self._on_trade: Optional[Callable] = None
|
|
self._on_order: Optional[Callable] = None
|
|
|
|
async def connect(self):
|
|
"""Connect to User Channel WebSocket."""
|
|
try:
|
|
self._running = True
|
|
|
|
logger.info(f"Connecting to User WebSocket at {WS_URL}...")
|
|
print(f" Connecting to {WS_URL}...")
|
|
|
|
async with ws_connect(
|
|
WS_URL,
|
|
**_ws_connect_kwargs(ping_interval=30, ping_timeout=10)
|
|
) as ws:
|
|
self._ws = ws
|
|
self._connected = True
|
|
logger.info("User WebSocket connected")
|
|
print(" WebSocket connection established")
|
|
|
|
# Subscribe to user channel with auth object
|
|
subscribe_msg = {
|
|
"type": "user",
|
|
"auth": {
|
|
"apiKey": self.api_key,
|
|
"secret": self.api_secret,
|
|
"passphrase": self.api_passphrase
|
|
}
|
|
}
|
|
await ws.send(json.dumps(subscribe_msg))
|
|
logger.info("Sent subscription message")
|
|
print(" Sent subscription message")
|
|
|
|
# Wait for response
|
|
try:
|
|
first_msg = await asyncio.wait_for(ws.recv(), timeout=5)
|
|
logger.info(f"First message: {first_msg[:200]}")
|
|
print(f" First response: {first_msg[:100]}...")
|
|
await self._process_message(first_msg)
|
|
except asyncio.TimeoutError:
|
|
logger.warning("No initial response from WebSocket")
|
|
print(" No initial response (timeout)")
|
|
|
|
# Listen for messages
|
|
while self._running:
|
|
try:
|
|
msg = await asyncio.wait_for(ws.recv(), timeout=60)
|
|
logger.debug(f"WS message: {msg[:100]}")
|
|
await self._process_message(msg)
|
|
except asyncio.TimeoutError:
|
|
# Send ping to keep connection alive
|
|
try:
|
|
await ws.ping()
|
|
except Exception:
|
|
break
|
|
except websockets.exceptions.ConnectionClosed:
|
|
logger.warning("User WebSocket connection closed")
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"Error receiving message: {e}")
|
|
break
|
|
|
|
self._connected = False
|
|
logger.info("User WebSocket disconnected")
|
|
print(" WebSocket disconnected")
|
|
|
|
except Exception as e:
|
|
logger.error(f"User WebSocket connection error: {e}")
|
|
print(f" Connection error: {e}")
|
|
self._connected = False
|
|
raise
|
|
|
|
async def _process_message(self, msg: str):
|
|
"""Process an incoming WebSocket message."""
|
|
try:
|
|
data = json.loads(msg)
|
|
except json.JSONDecodeError:
|
|
logger.warning(f"Invalid JSON: {msg[:100]}")
|
|
return
|
|
|
|
msg_type = data.get("type", "")
|
|
|
|
# Order updates
|
|
if msg_type in ("PLACEMENT", "UPDATE", "CANCELLATION"):
|
|
order_id = data.get("id", "")
|
|
if order_id and order_id in self._orders:
|
|
order = self._orders[order_id]
|
|
order.status = msg_type
|
|
order.size_matched = int(float(data.get("size_matched", 0)))
|
|
|
|
if msg_type == "MATCHED":
|
|
order.status = "MATCHED"
|
|
order.trades.append(data)
|
|
|
|
if self._on_order:
|
|
if asyncio.iscoroutinefunction(self._on_order):
|
|
await self._on_order(order)
|
|
else:
|
|
self._on_order(order)
|
|
|
|
# Trade updates
|
|
elif msg_type == "TRADE":
|
|
if self._on_trade:
|
|
if asyncio.iscoroutinefunction(self._on_trade):
|
|
await self._on_trade(data)
|
|
else:
|
|
self._on_trade(data)
|
|
|
|
# Initial subscription response
|
|
elif msg_type in ("subscribed", "OK", "ok"):
|
|
logger.info(f"Subscription confirmed: {msg[:150]}")
|
|
|
|
else:
|
|
logger.debug(f"Unhandled msg type={msg_type}: {msg[:100]}")
|
|
|
|
def register_order(self, order_id: str, asset_id: str, side: str, price: float, size: int):
|
|
"""Register an order we're tracking."""
|
|
self._orders[order_id] = OrderStatus(
|
|
order_id=order_id,
|
|
asset_id=asset_id,
|
|
side=side,
|
|
price=price,
|
|
original_size=size
|
|
)
|
|
|
|
def get_order_status(self, order_id: str) -> Optional[OrderStatus]:
|
|
"""Get current status for an order."""
|
|
return self._orders.get(order_id)
|
|
|
|
async def wait_for_order_match(self, order_id: str, timeout: float = 10.0) -> bool:
|
|
"""Wait until an order is matched or timeout."""
|
|
start = asyncio.get_event_loop().time()
|
|
while asyncio.get_event_loop().time() - start < timeout:
|
|
status = self._orders.get(order_id)
|
|
if status and (status.status == "MATCHED" or status.size_matched >= status.original_size):
|
|
return True
|
|
await asyncio.sleep(0.2)
|
|
return False
|
|
|
|
async def wait_for_order_place(self, order_id: str, timeout: float = 5.0) -> bool:
|
|
"""Wait until an order placement is confirmed or timeout."""
|
|
start = asyncio.get_event_loop().time()
|
|
while asyncio.get_event_loop().time() - start < timeout:
|
|
status = self._orders.get(order_id)
|
|
if status and status.status in ("PLACED", "UPDATE", "MATCHED"):
|
|
return True
|
|
await asyncio.sleep(0.15)
|
|
return False
|
|
|
|
@property
|
|
def is_connected(self) -> bool:
|
|
return self._connected and self._ws is not None
|
|
|
|
async def close(self):
|
|
"""Close the WebSocket connection."""
|
|
self._running = False
|
|
ws = self._ws
|
|
cm = self._ws_cm
|
|
self._ws = None
|
|
self._ws_cm = None
|
|
if ws and not getattr(ws, "closed", True):
|
|
try:
|
|
await ws.close()
|
|
except Exception:
|
|
pass
|
|
if cm is not None:
|
|
try:
|
|
await cm.__aexit__(None, None, None)
|
|
except Exception:
|
|
pass
|
|
self._connected = False |