From 0c60aa9b95fc9fb47ad8c8d8b289a72460bc506f Mon Sep 17 00:00:00 2001 From: nekrasovp Date: Fri, 11 Sep 2026 16:43:55 +0300 Subject: [PATCH] Fix graceful websocket reconnect and prepare 0.2.2 Refs #12. Restore subscriptions after EOF and make retry waits interruptible by close. --- CHANGELOG.md | 11 ++ docs/reconnects.md | 15 ++ llms-full.txt | 8 + llms.txt | 8 + pyproject.toml | 2 +- src/predxt/__init__.py | 2 +- src/predxt/kalshi/client.py | 119 +++++++------ src/predxt/opinion/client.py | 79 +++++---- src/predxt/polymarket/client.py | 122 +++++++------ src/predxt/utils/backoff.py | 16 ++ tests/test_graceful_reconnect.py | 287 +++++++++++++++++++++++++++++++ uv.lock | 2 +- 12 files changed, 540 insertions(+), 131 deletions(-) create mode 100644 tests/test_graceful_reconnect.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 83828d4..7c8f678 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,17 @@ All notable changes to `predxt` are documented here. +## 0.2.2 - 2026-09-11 + +- fixed graceful websocket stream exhaustion for Polymarket, Kalshi, and + Opinion: mark health disconnected, count the reconnect, wait for backoff, + and restore the connection and subscription +- stop Opinion's old heartbeat before reconnecting and restart it on success +- interrupt retry waits on explicit close and release websocket/heartbeat + resources when message iteration is cancelled or closed; repeated close + remains safe, and an explicit `connect()` can start a new session +- preserved public imports and the existing `VenueMessage.raw_data` contract + ## 0.2.1 - 2026-08-26 - fixed Polymarket, Kalshi, and Opinion connection manager shutdown so idle diff --git a/docs/reconnects.md b/docs/reconnects.md index 76e7868..3a51869 100644 --- a/docs/reconnects.md +++ b/docs/reconnects.md @@ -15,6 +15,21 @@ Connection retry behavior uses exponential backoff with jitter. Live systems should monitor `last_message_timestamp_ms`, reconnect counts, and process-level logs. +## Graceful stream endings + +Polymarket, Kalshi, and Opinion recover when the remote websocket iterator +ends normally as well as when it raises a connection error. Each disconnect +sets `connected` to false, increments `reconnect_count` once, and waits for +backoff before reconnecting and restoring the saved subscription. Opinion +stops its old heartbeat before waiting and starts a new one after connecting. +A completed socket is never repeatedly read without a retry delay. + +`close()` marks the client stopped before closing the socket and interrupts +pending retry waits. Cancelling the message reader or calling `aclose()` on +its iterator releases the connection and heartbeat without reconnecting. +Repeated `close()` calls are safe, including before the first connection. +After closing, call `connect()` explicitly to start a new session. + ## Connection manager shutdown Polymarket, Kalshi, and Opinion connection managers own their background diff --git a/llms-full.txt b/llms-full.txt index ae1087b..b6010f2 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -98,3 +98,11 @@ predxt stream polymarket --asset-id 1234567890 --limit 10 --jsonl - `docs/cli.md` - `docs/security.md` - `skills/predxt/SKILL.md` + +## Websocket recovery + +All three websocket clients recover from graceful stream endings with +backoff, updated connection health, and restored subscriptions. Opinion stops +its heartbeat during reconnect. Explicit `close()` interrupts retry waits; +cancelling or closing the message iterator releases the connection without +reconnecting. Call `connect()` explicitly to restart a closed client. diff --git a/llms.txt b/llms.txt index 38c947c..c9499ae 100644 --- a/llms.txt +++ b/llms.txt @@ -63,3 +63,11 @@ predxt stream polymarket --asset-id 1234567890 --limit 10 --jsonl - Prefer `OrderBookState` for simple dashboards; do not claim execution-grade semantics. More context: `llms-full.txt`, `docs/`, `skills/predxt/SKILL.md`. + +## Websocket recovery + +All three websocket clients recover from graceful stream endings with +backoff, updated connection health, and restored subscriptions. Opinion stops +its heartbeat during reconnect. Explicit `close()` interrupts retry waits; +cancelling or closing the message iterator releases the connection without +reconnecting. Call `connect()` explicitly to restart a closed client. diff --git a/pyproject.toml b/pyproject.toml index 4612179..93d7306 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "predxt" -version = "0.2.1" +version = "0.2.2" description = "Read-only prediction market data clients for Polymarket, Kalshi, and Opinion" readme = "README.md" requires-python = ">=3.12" diff --git a/src/predxt/__init__.py b/src/predxt/__init__.py index 9952f82..7d65464 100644 --- a/src/predxt/__init__.py +++ b/src/predxt/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.2.1" +__version__ = "0.2.2" from .base import BaseWsClient, HealthMetrics, VenueMessage, build_venue_message from .events import ( diff --git a/src/predxt/kalshi/client.py b/src/predxt/kalshi/client.py index 9e40158..7ffee7e 100644 --- a/src/predxt/kalshi/client.py +++ b/src/predxt/kalshi/client.py @@ -18,7 +18,7 @@ ) from predxt.kalshi.auth import build_kalshi_auth_headers from predxt.kalshi.parser import parse_message -from predxt.utils.backoff import ExponentialBackoff +from predxt.utils.backoff import ExponentialBackoff, _wait_for_backoff logger = logging.getLogger(__name__) @@ -35,6 +35,7 @@ def __init__( self._ws_url = ws_url self._ws: Any | None = None self._connected = False + self._closed = asyncio.Event() self._connect_time: Optional[float] = None self._channels: list[str] = [] self._params: dict[str, Any] = {} @@ -44,17 +45,22 @@ def __init__( self._max_reconnect_attempts = max_reconnect_attempts async def connect(self, auth_params: Optional[dict[str, Any]] = None) -> None: + self._closed.clear() if auth_params is not None: self._auth_params = auth_params ws_path = urlparse(self._ws_url).path or "/trade-api/ws/v2" headers = self._build_auth_headers(self._auth_params, ws_path=ws_path) attempts = 0 - while attempts < self._max_reconnect_attempts: + while not self._closed.is_set() and attempts < self._max_reconnect_attempts: try: - self._ws = await websockets.connect( + ws = await websockets.connect( self._ws_url, additional_headers=headers ) + if self._closed.is_set(): + await ws.close() + return + self._ws = ws self._connected = True self._connect_time = time.time() self._backoff.reset() @@ -66,8 +72,10 @@ async def connect(self, auth_params: Optional[dict[str, Any]] = None) -> None: attempts += 1 self._health.last_error = str(exc) self._health.last_error_timestamp_ms = time.time() * 1000 - await asyncio.sleep(self._backoff.next_delay()) + await _wait_for_backoff(self._backoff, self._closed) + if self._closed.is_set(): + return raise ConnectionError( f"Failed to connect to Kalshi after {self._max_reconnect_attempts} attempts" ) @@ -104,60 +112,73 @@ async def unsubscribe(self, channels: list[str]) -> None: self._health.messages_sent += 1 async def close(self) -> None: - if self._ws: - await self._ws.close() + self._closed.set() self._connected = False self._health.connected = False + ws, self._ws = self._ws, None + if ws is not None: + await ws.close() async def messages(self) -> AsyncIterator[VenueMessage]: seen_hashes: set[str] = set() - while True: - try: - if not self._connected: - await self.connect(self._auth_params) - if self._channels: - await self.subscribe(self._channels, self._params) - - if self._ws is None: - continue - async for raw_msg in self._ws: - payload = json.loads(raw_msg) - parsed = parse_message(payload) - if not parsed: - continue - - dedupe_hash = parsed.get("hash") - if not dedupe_hash: - dedupe_hash = hashlib.sha1( - json.dumps(parsed, sort_keys=True).encode("utf-8") - ).hexdigest() - - if dedupe_hash in seen_hashes: - continue - seen_hashes.add(dedupe_hash) - - vm = build_venue_message( - venue="kalshi", - raw_data=parsed, - timestamp_ms=time.time() * 1000, - ) - - self._health.messages_received += 1 - self._health.last_message_timestamp_ms = vm.timestamp_ms - yield vm - - except (websockets.exceptions.ConnectionClosed, OSError): + try: + while not self._closed.is_set(): + try: + if not self._connected: + await self.connect(self._auth_params) + if self._closed.is_set(): + return + if self._channels: + await self.subscribe(self._channels, self._params) + + if self._closed.is_set(): + return + if self._ws is None: + raise OSError("Websocket unavailable") + async for raw_msg in self._ws: + payload = json.loads(raw_msg) + parsed = parse_message(payload) + if not parsed: + continue + + dedupe_hash = parsed.get("hash") + if not dedupe_hash: + dedupe_hash = hashlib.sha1( + json.dumps(parsed, sort_keys=True).encode("utf-8") + ).hexdigest() + + if dedupe_hash in seen_hashes: + continue + seen_hashes.add(dedupe_hash) + + vm = build_venue_message( + venue="kalshi", + raw_data=parsed, + timestamp_ms=time.time() * 1000, + ) + + self._health.messages_received += 1 + self._health.last_message_timestamp_ms = vm.timestamp_ms + yield vm + + except (websockets.exceptions.ConnectionClosed, OSError): + pass + except Exception: + logger.exception("Unexpected Kalshi websocket error") + raise + + # Both graceful EOF and transport errors require a paced reconnect. + if self._closed.is_set(): + return self._connected = False self._health.connected = False self._health.reconnect_count += 1 - await asyncio.sleep(self._backoff.next_delay()) - except Exception: - logger.exception("Unexpected Kalshi websocket error") - raise - - if len(seen_hashes) > 10000: - seen_hashes.clear() + await _wait_for_backoff(self._backoff, self._closed) + if len(seen_hashes) > 10000: + seen_hashes.clear() + finally: + await self.close() def health(self) -> HealthMetrics: self._health.connected = self._connected diff --git a/src/predxt/opinion/client.py b/src/predxt/opinion/client.py index 8d70c6f..392f8fc 100644 --- a/src/predxt/opinion/client.py +++ b/src/predxt/opinion/client.py @@ -16,7 +16,7 @@ build_venue_message, ) from predxt.opinion.parser import parse_message -from predxt.utils.backoff import ExponentialBackoff +from predxt.utils.backoff import ExponentialBackoff, _wait_for_backoff logger = logging.getLogger(__name__) @@ -38,6 +38,7 @@ def __init__( self._max_reconnect_attempts = max_reconnect_attempts self._ws: Any | None = None self._connected = False + self._closed = asyncio.Event() self._connect_time: Optional[float] = None self._channels: list[str] = [] self._params: dict[str, Any] = {} @@ -47,6 +48,7 @@ def __init__( self._heartbeat_task: asyncio.Task | None = None async def connect(self, auth_params: Optional[dict[str, Any]] = None) -> None: + self._closed.clear() if auth_params is not None: self._auth_params = auth_params api_key = self._resolve_api_key() @@ -54,9 +56,13 @@ async def connect(self, auth_params: Optional[dict[str, Any]] = None) -> None: raise ValueError("Opinion websocket requires an API key") attempts = 0 - while attempts < self._max_reconnect_attempts: + while not self._closed.is_set() and attempts < self._max_reconnect_attempts: try: - self._ws = await websockets.connect(self._authenticated_url(api_key)) + ws = await websockets.connect(self._authenticated_url(api_key)) + if self._closed.is_set(): + await ws.close() + return + self._ws = ws self._connected = True self._connect_time = time.time() self._backoff.reset() @@ -69,8 +75,10 @@ async def connect(self, auth_params: Optional[dict[str, Any]] = None) -> None: attempts += 1 self._health.last_error = str(exc) self._health.last_error_timestamp_ms = time.time() * 1000 - await asyncio.sleep(self._backoff.next_delay()) + await _wait_for_backoff(self._backoff, self._closed) + if self._closed.is_set(): + return raise ConnectionError( f"Failed to connect to Opinion after {self._max_reconnect_attempts} attempts" ) @@ -96,41 +104,54 @@ async def unsubscribe(self, channels: list[str]) -> None: self._health.messages_sent += 1 async def close(self) -> None: - await self._stop_heartbeat() - if self._ws: - await self._ws.close() + self._closed.set() self._connected = False self._health.connected = False + ws, self._ws = self._ws, None + await self._stop_heartbeat() + if ws is not None: + await ws.close() async def messages(self) -> AsyncIterator[VenueMessage]: seen_hashes: set[str] = set() - while True: - try: - if not self._connected: - await self.connect(self._auth_params) - if self._channels: - await self.subscribe(self._channels, self._params) + try: + while not self._closed.is_set(): + try: + if not self._connected: + await self.connect(self._auth_params) + if self._closed.is_set(): + return + if self._channels: + await self.subscribe(self._channels, self._params) + + if self._closed.is_set(): + return + if self._ws is None: + raise OSError("Websocket unavailable") + async for raw_msg in self._ws: + async for vm in self._handle_raw_message( + raw_msg, + seen_hashes=seen_hashes, + ): + yield vm + except (websockets.exceptions.ConnectionClosed, OSError): + pass + except Exception: + logger.exception("Unexpected Opinion websocket error") + raise - if self._ws is None: - continue - async for raw_msg in self._ws: - async for vm in self._handle_raw_message( - raw_msg, - seen_hashes=seen_hashes, - ): - yield vm - except (websockets.exceptions.ConnectionClosed, OSError): + # Both graceful EOF and transport errors require a paced reconnect. + if self._closed.is_set(): + return self._connected = False self._health.connected = False self._health.reconnect_count += 1 await self._stop_heartbeat() - await asyncio.sleep(self._backoff.next_delay()) - except Exception: - logger.exception("Unexpected Opinion websocket error") - raise - - if len(seen_hashes) > 10000: - seen_hashes.clear() + await _wait_for_backoff(self._backoff, self._closed) + if len(seen_hashes) > 10000: + seen_hashes.clear() + finally: + await self.close() def health(self) -> HealthMetrics: self._health.connected = self._connected diff --git a/src/predxt/polymarket/client.py b/src/predxt/polymarket/client.py index 6c225be..5346a50 100644 --- a/src/predxt/polymarket/client.py +++ b/src/predxt/polymarket/client.py @@ -14,7 +14,7 @@ VenueMessage, build_venue_message, ) -from predxt.utils.backoff import ExponentialBackoff +from predxt.utils.backoff import ExponentialBackoff, _wait_for_backoff from .parser import parse_message @@ -28,6 +28,7 @@ class PolymarketWsClient(BaseWsClient): def __init__(self, max_reconnect_attempts: int = MAX_RECONNECT_ATTEMPTS): self._ws: Any | None = None self._connected = False + self._closed = asyncio.Event() self._backoff = ExponentialBackoff(base_seconds=1, max_seconds=60) self._health = HealthMetrics(connected=False) self._channels: list[str] = [] @@ -41,12 +42,20 @@ async def connect(self, auth_params: Optional[dict[str, Any]] = None) -> None: Sets _connect_time, clears last_error on success and records last_error on failure. """ + self._closed.clear() if auth_params: logger.warning("Polymarket WS is public; auth_params ignored") - while self._reconnect_attempts < self._max_reconnect_attempts: + while ( + not self._closed.is_set() + and self._reconnect_attempts < self._max_reconnect_attempts + ): try: - self._ws = await websockets.connect(self.WS_URL) + ws = await websockets.connect(self.WS_URL) + if self._closed.is_set(): + await ws.close() + return + self._ws = ws self._connected = True # set connection timestamp used by health(). self._connect_time = time.time() @@ -65,8 +74,10 @@ async def connect(self, auth_params: Optional[dict[str, Any]] = None) -> None: self._health.last_error = str(e) self._health.last_error_timestamp_ms = time.time() * 1000 logger.warning(f"Connection failed, retry in {delay}s: {e}") - await asyncio.sleep(delay) + await _wait_for_backoff(self._backoff, self._closed, delay=delay) + if self._closed.is_set(): + return raise ConnectionError( f"Failed to connect after {self._max_reconnect_attempts} attempts" ) @@ -118,11 +129,12 @@ async def unsubscribe(self, channels: list[str]) -> None: self._connected = False async def close(self) -> None: - if self._ws: - await self._ws.close() - self._connected = False - self._health.connected = False - logger.info("Closed Polymarket WS") + self._closed.set() + self._connected = False + self._health.connected = False + ws, self._ws = self._ws, None + if ws is not None: + await ws.close() async def messages(self) -> AsyncIterator[VenueMessage]: """Stream parsed Polymarket messages with deduplication and health updates. @@ -131,60 +143,70 @@ async def messages(self) -> AsyncIterator[VenueMessage]: """ seen_hashes: set[str] = set() - while True: - try: - if not self._connected: - await self.connect() - if self._channels: - await self.subscribe(self._channels, self._params) - - # websockets client returns an async iterator; support AsyncMock whose __aiter__ may be an async generator - # Prefer iterating directly over the websocket object when possible. + try: + while not self._closed.is_set(): try: - # Prefer standard async iteration - ws = self._ws - if ws is None: - continue - async for raw_msg in ws: - async for vm in self._handle_raw_message( - raw_msg, seen_hashes=seen_hashes - ): - yield vm - except TypeError: - # The websocket mock may expose __aiter__ as an async generator object - # that isn't directly iterable; attempt to call it and iterate the result. + if not self._connected: + await self.connect() + if self._closed.is_set(): + return + if self._channels: + await self.subscribe(self._channels, self._params) + + # websockets client returns an async iterator; support AsyncMock whose __aiter__ may be an async generator + # Prefer iterating directly over the websocket object when possible. try: + # Prefer standard async iteration ws = self._ws + if self._closed.is_set(): + return if ws is None: - continue - aiter_obj = ws.__aiter__() - # If __aiter__() returned a coroutine that yields an async generator, await it - if asyncio.iscoroutine(aiter_obj): - aiter_obj = await aiter_obj - - async for raw_msg in aiter_obj: + raise OSError("Websocket unavailable") + async for raw_msg in ws: async for vm in self._handle_raw_message( raw_msg, seen_hashes=seen_hashes ): yield vm - except Exception as e: - logger.error(f"Failed iterating websocket mock: {e}") - # end of mock iterator handling - - except (websockets.exceptions.ConnectionClosed, OSError) as e: + except TypeError: + # The websocket mock may expose __aiter__ as an async generator object + # that isn't directly iterable; attempt to call it and iterate the result. + try: + ws = self._ws + if self._closed.is_set(): + return + if ws is None: + raise OSError("Websocket unavailable") + aiter_obj = ws.__aiter__() + # If __aiter__() returned a coroutine that yields an async generator, await it + if asyncio.iscoroutine(aiter_obj): + aiter_obj = await aiter_obj + + async for raw_msg in aiter_obj: + async for vm in self._handle_raw_message( + raw_msg, seen_hashes=seen_hashes + ): + yield vm + except Exception as e: + logger.error(f"Failed iterating websocket mock: {e}") + # end of mock iterator handling + + except (websockets.exceptions.ConnectionClosed, OSError) as e: + logger.warning(f"Connection lost: {e}, reconnecting...") + except Exception as e: + logger.error(f"Unexpected error: {e}") + raise + + # Both graceful EOF and transport errors require a paced reconnect. + if self._closed.is_set(): + return self._connected = False self._health.connected = False self._health.reconnect_count += 1 - logger.warning(f"Connection lost: {e}, reconnecting...") - await asyncio.sleep(self._backoff.next_delay()) - - # On reconnect, preserve seen_hashes for a short time but don't let it grow unbounded + await _wait_for_backoff(self._backoff, self._closed) if len(seen_hashes) > 10000: seen_hashes.clear() - - except Exception as e: - logger.error(f"Unexpected error: {e}") - raise + finally: + await self.close() def health(self) -> HealthMetrics: self._health.connected = self._connected diff --git a/src/predxt/utils/backoff.py b/src/predxt/utils/backoff.py index fea31b9..73606d8 100644 --- a/src/predxt/utils/backoff.py +++ b/src/predxt/utils/backoff.py @@ -1,3 +1,4 @@ +import asyncio import random @@ -15,3 +16,18 @@ def next_delay(self) -> float: def reset(self) -> None: self._attempt = 0 + + +async def _wait_for_backoff( + backoff: ExponentialBackoff, + closed: asyncio.Event, + *, + delay: float | None = None, +) -> None: + """Pace retries while allowing an explicit close to interrupt the wait.""" + if delay is None: + delay = backoff.next_delay() + try: + await asyncio.wait_for(closed.wait(), timeout=delay) + except TimeoutError: + pass diff --git a/tests/test_graceful_reconnect.py b/tests/test_graceful_reconnect.py new file mode 100644 index 0000000..72c9d28 --- /dev/null +++ b/tests/test_graceful_reconnect.py @@ -0,0 +1,287 @@ +"""Offline lifecycle regressions; no venue credentials or network calls are used.""" + +import asyncio +import importlib +import json +from collections import deque +from unittest.mock import AsyncMock, Mock + +import pytest +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK +from websockets.frames import Close + +from predxt.kalshi import KalshiWsClient +from predxt.kalshi.connection_manager import KalshiWsConnectionManager +from predxt.opinion import OpinionWsClient +from predxt.opinion.connection_manager import OpinionWsConnectionManager +from predxt.polymarket import PolymarketWsClient +from predxt.polymarket.connection_manager import WsConnectionManager + + +class Socket: + def __init__(self, *payloads, idle=False, error=None): + self.payloads = deque(json.dumps(p) for p in payloads) + self.idle = idle + self.error = error + self.iterations = 0 + self.reading = asyncio.Event() + self.closed = asyncio.Event() + self.send = AsyncMock() + self.close = AsyncMock(side_effect=self.closed.set) + + def __aiter__(self): + self.iterations += 1 + # Fail synchronously on the old bug, even if it starves the event loop. + assert self.iterations == 1, "exhausted socket iterated again (busy loop)" + return self + + async def __anext__(self): + self.reading.set() + if self.payloads: + return self.payloads.popleft() + if self.idle: + await self.closed.wait() + if self.error is not None: + raise self.error + raise StopAsyncIteration + + +@pytest.fixture(params=["polymarket", "kalshi", "opinion"]) +def venue(request, monkeypatch): + name = request.param + if name == "polymarket": + client = PolymarketWsClient() + auth = None + channels, params = ["market"], {"assets_ids": ["test-asset"]} + payload = { + "event_type": "price_change", + "asset_id": "test-asset", + "price": "0.55", + "hash": "test-hash", + } + manager = WsConnectionManager(client=client, initial_assets=["test-asset"]) + elif name == "kalshi": + client = KalshiWsClient() + # Stub signing itself; no precomputed signatures or credentials are needed. + monkeypatch.setattr(client, "_build_auth_headers", lambda *a, **kw: {}) + auth = {"test_auth_marker": "preserved"} + channels, params = ["orderbook_delta"], {"market_tickers": ["TEST-MARKET"]} + payload = { + "type": "orderbook_delta", + "msg": { + "market_ticker": "TEST-MARKET", + "price": "0.55", + "hash": "test-hash", + }, + } + manager = KalshiWsConnectionManager( + client=client, initial_markets=["TEST-MARKET"] + ) + else: + client = OpinionWsClient(heartbeat_seconds=3600) + monkeypatch.setattr(client, "_resolve_api_key", lambda: "offline-test-only") + auth = {"test_auth_marker": "preserved"} + channels, params = ["market.depth.diff"], {"market_ids": ["123"]} + payload = { + "msgType": "market.depth.diff", + "marketId": 123, + "price": "0.55", + "size": "12", + } + manager = OpinionWsConnectionManager(client=client, initial_markets=["123"]) + return name, client, auth, channels, params, payload, manager + + +async def start_client(venue, monkeypatch, *sockets): + name, client, auth, channels, params, _, _ = venue + connect = AsyncMock(side_effect=sockets) + monkeypatch.setattr("websockets.connect", connect) + await client.connect(auth) + await client.subscribe(channels, params) + return connect + + +async def bounded(awaitable): + return await asyncio.wait_for(awaitable, timeout=1) + + +@pytest.mark.parametrize("ending", ["eof", "oserror", "close_ok", "close_error"]) +async def test_stream_end_backoff_reconnect_resubscribe(venue, monkeypatch, ending): + name, client, auth, _, _, payload, _ = venue + errors = { + "eof": None, + "oserror": OSError("offline disconnect"), + "close_ok": ConnectionClosedOK(Close(1000, ""), Close(1000, ""), True), + "close_error": ConnectionClosedError(Close(1011, ""), None, None), + } + first, second = Socket(error=errors[ending]), Socket(payload, idle=True) + connect = await start_client(venue, monkeypatch, first, second) + health = ( + client.health() + ) # Hold the metrics object: no refresh may mask stale state. + heartbeat = getattr(client, "_heartbeat_task", None) + entered, release = asyncio.Event(), asyncio.Event() + delays = [] + + async def wait_backoff(backoff, closed): + delays.append(backoff.next_delay()) + assert client._connected is False + assert health.connected is False + assert health.reconnect_count == 1 + if heartbeat is not None: + assert heartbeat.done() + assert client._heartbeat_task is None + entered.set() + await release.wait() + + module = importlib.import_module(f"predxt.{name}.client") + monkeypatch.setattr(module, "_wait_for_backoff", wait_backoff, raising=False) + stream = client.messages() + reader = asyncio.create_task(anext(stream)) + try: + await bounded(entered.wait()) + # Give other tasks turns; retry must remain suspended behind backoff. + for _ in range(3): + await asyncio.sleep(0) + assert connect.await_count == 1 + assert first.iterations == 1 + assert not reader.done() + release.set() + message = await bounded(reader) + assert len(delays) == 1 and delays[0] >= 1 + assert connect.await_count == 2 + assert second.iterations == 1 + assert first.send.await_count == second.send.await_count == 1 + sent_first = json.loads(first.send.call_args.args[0]) + sent_second = json.loads(second.send.call_args.args[0]) + if name == "kalshi": + sent_first.pop("id") + sent_second.pop("id") + assert sent_first == sent_second + assert health.connected is True and client._connected is True + assert health.reconnect_count == 1 + assert health.messages_received == 1 + assert message.venue == name + assert message.raw_data["price"] == "0.55" + assert message.received_at_ms == message.timestamp_ms + if name == "kalshi": + assert message.raw_data["raw"] == payload + if auth is not None: + assert client._auth_params == auth + if heartbeat is not None: + assert client._heartbeat_task is not heartbeat + assert not client._heartbeat_task.done() + finally: + reader.cancel() + await asyncio.gather(reader, return_exceptions=True) + await bounded(stream.aclose()) + await bounded(client.close()) + assert not health.connected + assert connect.await_count == 2 + + +@pytest.mark.parametrize("action", ["close", "cancel", "stop"]) +@pytest.mark.parametrize("phase", ["idle", "backoff"]) +async def test_shutdown_is_prompt_and_does_not_reconnect( + venue, monkeypatch, phase, action +): + name, client, auth, _, _, _, manager = venue + socket = Socket(idle=phase == "idle") + connect = AsyncMock(return_value=socket) + monkeypatch.setattr("websockets.connect", connect) + backoff_entered = asyncio.Event() + + def next_delay(): + backoff_entered.set() + return 3600 # A real wait that must be interrupted, never slept through. + + monkeypatch.setattr(client._backoff, "next_delay", Mock(side_effect=next_delay)) + stream = None + if action == "stop": + if name == "polymarket": + await manager.start() + else: + await manager.start(auth) + reader = None + else: + await client.connect(auth) + stream = client.messages() + reader = asyncio.create_task(anext(stream)) + try: + await bounded((socket.reading if phase == "idle" else backoff_entered).wait()) + if action == "stop": + await bounded(manager.stop()) + await bounded(manager.stop()) + elif action == "close": + await bounded(client.close()) + await bounded(client.close()) + with pytest.raises(StopAsyncIteration): + await bounded(reader) + else: + reader.cancel() + with pytest.raises(asyncio.CancelledError): + await bounded(reader) + await bounded(client.close()) + await bounded(client.close()) + assert not client._connected + assert not client.health().connected + assert client.health().reconnect_count == (1 if phase == "backoff" else 0) + assert connect.await_count == 1 + assert socket.close.await_count == 1 + assert socket.iterations == 1 + assert getattr(client, "_heartbeat_task", None) is None + finally: + if reader is not None: + reader.cancel() + await asyncio.gather(reader, return_exceptions=True) + if stream is not None: + await stream.aclose() + await client.close() + + +async def test_close_before_messages_and_explicit_restart(venue, monkeypatch): + _, client, auth, _, _, payload, _ = venue + socket = Socket(payload, idle=True) + connect = AsyncMock(return_value=socket) + monkeypatch.setattr("websockets.connect", connect) + await bounded(client.close()) + await bounded(client.close()) + with pytest.raises(StopAsyncIteration): + await bounded(anext(client.messages())) + assert connect.await_count == 0 + await client.connect(auth) + stream = client.messages() + try: + message = await bounded(anext(stream)) + assert message.raw_data["price"] == "0.55" + finally: + await stream.aclose() + assert connect.await_count == 1 + assert not client.health().connected + + +async def test_close_during_connection_does_not_restore_health(venue, monkeypatch): + _, client, auth, _, _, _, _ = venue + entered, release = asyncio.Event(), asyncio.Event() + socket = Socket() + + async def connect_socket(*args, **kwargs): + entered.set() + await release.wait() + return socket + + connect = AsyncMock(side_effect=connect_socket) + monkeypatch.setattr("websockets.connect", connect) + task = asyncio.create_task(client.connect(auth)) + try: + await bounded(entered.wait()) + await bounded(client.close()) + release.set() + await bounded(task) + assert not client.health().connected + assert connect.await_count == socket.close.await_count == 1 + assert getattr(client, "_heartbeat_task", None) is None + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await client.close() diff --git a/uv.lock b/uv.lock index 91a8060..4e52557 100644 --- a/uv.lock +++ b/uv.lock @@ -850,7 +850,7 @@ wheels = [ [[package]] name = "predxt" -version = "0.2.1" +version = "0.2.2" source = { editable = "." } dependencies = [ { name = "cryptography" },