From c1118162ee919f8faa9ab3b19a2627148c9c1199 Mon Sep 17 00:00:00 2001 From: nekrasovp Date: Wed, 26 Aug 2026 01:35:11 +0300 Subject: [PATCH] Fix idle connection manager shutdown Cancel and await idle websocket readers across all venue managers while allowing in-flight dispatch to complete. Add lifecycle regressions and document public stop semantics.\n\nCloses #9. --- CHANGELOG.md | 7 ++ README.md | 5 + docs/reconnects.md | 18 +++ llms-full.txt | 7 ++ llms.txt | 2 + src/predxt/kalshi/connection_manager.py | 19 +++- src/predxt/opinion/connection_manager.py | 19 +++- src/predxt/polymarket/connection_manager.py | 19 +++- tests/test_connection_manager_lifecycle.py | 119 ++++++++++++++++++++ 9 files changed, 206 insertions(+), 9 deletions(-) create mode 100644 tests/test_connection_manager_lifecycle.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c7f766..35217a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,13 @@ All notable changes to `predxt` are documented here. +## Unreleased + +- fixed Polymarket, Kalshi, and Opinion connection manager shutdown so idle + websocket streams cannot block `stop()` +- made pre-start and repeated `stop()` calls safe without requiring consumers + to cancel private message tasks + ## 0.2.0 - 2026-06-05 - added read-only REST client base, normalized market/orderbook models, and diff --git a/README.md b/README.md index 0c823a4..305a967 100644 --- a/README.md +++ b/README.md @@ -124,6 +124,11 @@ book = await client.get_orderbook("CLOB_TOKEN_ID") await client.close() ``` +Connection managers own their background websocket reader. Use their public +`start()` and `stop()` methods for lifecycle management; `stop()` also handles +idle streams and is safe to call before `start()` or more than once. Consumers +do not need to access private task attributes. + ## CLI Offline parser demo: diff --git a/docs/reconnects.md b/docs/reconnects.md index 8fc63a0..76e7868 100644 --- a/docs/reconnects.md +++ b/docs/reconnects.md @@ -14,3 +14,21 @@ Clients expose `health()`: Connection retry behavior uses exponential backoff with jitter. Live systems should monitor `last_message_timestamp_ms`, reconnect counts, and process-level logs. + +## Connection manager shutdown + +Polymarket, Kalshi, and Opinion connection managers own their background +message reader. Always shut a manager down through its public lifecycle: + +```python +await manager.start() +try: + ... +finally: + await manager.stop() +``` + +`stop()` cancels and awaits the message reader before closing the websocket +client, so it completes even when the stream is idle. Calling `stop()` before +`start()` or more than once is safe. Consumers should not inspect or cancel +private task attributes. diff --git a/llms-full.txt b/llms-full.txt index 7c4dd7d..ae1087b 100644 --- a/llms-full.txt +++ b/llms-full.txt @@ -60,6 +60,13 @@ Shared REST classes include `MarketSummary`, `MarketDetail`, `OrderBookSnapshot`, `OrderBookLevel`, `VenueCredentialStatus`, and `VenueApiError`. +### Connection manager lifecycle + +Polymarket, Kalshi, and Opinion connection managers own their background +message reader. Use the public `start()` and `stop()` methods. `stop()` handles +idle streams and is safe before `start()` or when called repeatedly; consumers +must not inspect or cancel private message task attributes. + ### CLI ```bash diff --git a/llms.txt b/llms.txt index 0be1ece..38c947c 100644 --- a/llms.txt +++ b/llms.txt @@ -58,6 +58,8 @@ predxt stream polymarket --asset-id 1234567890 --limit 10 --jsonl - Preserve raw payload access. - Keep generated apps read-only. - Use env vars or secret managers for Kalshi and Opinion credentials. +- Use connection managers' public `start()` and `stop()` lifecycle; never + inspect or cancel their private message tasks. - Prefer `OrderBookState` for simple dashboards; do not claim execution-grade semantics. More context: `llms-full.txt`, `docs/`, `skills/predxt/SKILL.md`. diff --git a/src/predxt/kalshi/connection_manager.py b/src/predxt/kalshi/connection_manager.py index 8312963..ab3ffaa 100644 --- a/src/predxt/kalshi/connection_manager.py +++ b/src/predxt/kalshi/connection_manager.py @@ -62,6 +62,8 @@ def __init__( self._market_callbacks: Dict[str, List[MessageCallback]] = defaultdict(list) self._messages_task: Optional[asyncio.Task] = None self._stop_event = asyncio.Event() + self._dispatch_complete = asyncio.Event() + self._dispatch_complete.set() self._lock = asyncio.Lock() def add_market(self, market_ticker: str) -> None: @@ -92,8 +94,15 @@ async def start(self, auth_params: Optional[dict] = None) -> None: async def stop(self) -> None: self._stop_event.set() - if self._messages_task: - await asyncio.gather(self._messages_task, return_exceptions=True) + messages_task = self._messages_task + if messages_task and not self._dispatch_complete.is_set(): + await self._dispatch_complete.wait() + if messages_task and not messages_task.done(): + messages_task.cancel() + if messages_task: + await asyncio.gather(messages_task, return_exceptions=True) + if self._messages_task is messages_task: + self._messages_task = None await self._client.close() async def refresh_subscription(self) -> None: @@ -109,7 +118,11 @@ def config(self) -> KalshiSubscriptionConfig: async def _message_loop(self) -> None: try: async for msg in self._client.messages(): - await self._dispatch(msg) + self._dispatch_complete.clear() + try: + await self._dispatch(msg) + finally: + self._dispatch_complete.set() if self._stop_event.is_set(): break except asyncio.CancelledError: diff --git a/src/predxt/opinion/connection_manager.py b/src/predxt/opinion/connection_manager.py index 9d47624..c26fad3 100644 --- a/src/predxt/opinion/connection_manager.py +++ b/src/predxt/opinion/connection_manager.py @@ -54,6 +54,8 @@ def __init__( self._market_callbacks: Dict[str, List[MessageCallback]] = defaultdict(list) self._messages_task: Optional[asyncio.Task] = None self._stop_event = asyncio.Event() + self._dispatch_complete = asyncio.Event() + self._dispatch_complete.set() self._lock = asyncio.Lock() def add_market(self, market_id: str) -> None: @@ -86,8 +88,15 @@ async def start(self, auth_params: Optional[dict] = None) -> None: async def stop(self) -> None: self._stop_event.set() - if self._messages_task: - await asyncio.gather(self._messages_task, return_exceptions=True) + messages_task = self._messages_task + if messages_task and not self._dispatch_complete.is_set(): + await self._dispatch_complete.wait() + if messages_task and not messages_task.done(): + messages_task.cancel() + if messages_task: + await asyncio.gather(messages_task, return_exceptions=True) + if self._messages_task is messages_task: + self._messages_task = None await self._client.close() async def refresh_subscription(self) -> None: @@ -103,7 +112,11 @@ def config(self) -> OpinionSubscriptionConfig: async def _message_loop(self) -> None: try: async for msg in self._client.messages(): - await self._dispatch(msg) + self._dispatch_complete.clear() + try: + await self._dispatch(msg) + finally: + self._dispatch_complete.set() if self._stop_event.is_set(): break except asyncio.CancelledError: diff --git a/src/predxt/polymarket/connection_manager.py b/src/predxt/polymarket/connection_manager.py index b2b0461..099e48c 100644 --- a/src/predxt/polymarket/connection_manager.py +++ b/src/predxt/polymarket/connection_manager.py @@ -74,6 +74,8 @@ def __init__( self._asset_callbacks: Dict[str, List[MessageCallback]] = defaultdict(list) self._messages_task: Optional[asyncio.Task] = None self._stop_event = asyncio.Event() + self._dispatch_complete = asyncio.Event() + self._dispatch_complete.set() self._lock = asyncio.Lock() # ------------------------------------------------------------------ @@ -107,8 +109,15 @@ async def start(self) -> None: async def stop(self) -> None: self._stop_event.set() - if self._messages_task: - await asyncio.gather(self._messages_task, return_exceptions=True) + messages_task = self._messages_task + if messages_task and not self._dispatch_complete.is_set(): + await self._dispatch_complete.wait() + if messages_task and not messages_task.done(): + messages_task.cancel() + if messages_task: + await asyncio.gather(messages_task, return_exceptions=True) + if self._messages_task is messages_task: + self._messages_task = None await self._client.close() async def refresh_subscription(self) -> None: @@ -124,7 +133,11 @@ def config(self) -> SubscriptionConfig: async def _message_loop(self) -> None: try: async for msg in self._client.messages(): - await self._dispatch(msg) + self._dispatch_complete.clear() + try: + await self._dispatch(msg) + finally: + self._dispatch_complete.set() if self._stop_event.is_set(): break except asyncio.CancelledError: diff --git a/tests/test_connection_manager_lifecycle.py b/tests/test_connection_manager_lifecycle.py new file mode 100644 index 0000000..613fa36 --- /dev/null +++ b/tests/test_connection_manager_lifecycle.py @@ -0,0 +1,119 @@ +import asyncio +import logging + +import pytest + +from predxt.base import VenueMessage +from predxt.kalshi.connection_manager import ( + KalshiSubscriptionConfig, + KalshiWsConnectionManager, +) +from predxt.opinion.connection_manager import ( + OpinionSubscriptionConfig, + OpinionWsConnectionManager, +) +from predxt.polymarket.connection_manager import ( + SubscriptionConfig, + WsConnectionManager, +) + + +class IdleClient: + def __init__(self) -> None: + self.messages_started = asyncio.Event() + self.messages_stopped = asyncio.Event() + self._release_messages = asyncio.Event() + self.closed = False + self.close_calls = 0 + + async def connect(self, *args, **kwargs) -> None: + pass + + async def subscribe(self, channels, params) -> None: + pass + + async def close(self) -> None: + self.close_calls += 1 + self.closed = True + + async def messages(self): + self.messages_started.set() + try: + await self._release_messages.wait() + yield VenueMessage() + finally: + self.messages_stopped.set() + + +MANAGER_FACTORIES = [ + pytest.param( + lambda client: WsConnectionManager( + client=client, + config=SubscriptionConfig(assets_ids={"asset-1"}), + ), + id="polymarket", + ), + pytest.param( + lambda client: KalshiWsConnectionManager( + client=client, + config=KalshiSubscriptionConfig(market_tickers={"market-1"}), + ), + id="kalshi", + ), + pytest.param( + lambda client: OpinionWsConnectionManager( + client=client, + config=OpinionSubscriptionConfig(market_ids={"market-1"}), + ), + id="opinion", + ), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("manager_factory", MANAGER_FACTORIES) +async def test_stop_completes_for_idle_stream(manager_factory, caplog): + client = IdleClient() + manager = manager_factory(client) + + await manager.start() + await asyncio.wait_for(client.messages_started.wait(), timeout=0.1) + + with caplog.at_level(logging.ERROR): + await asyncio.wait_for(manager.stop(), timeout=0.1) + + assert client.closed is True + assert client.close_calls == 1 + assert client.messages_stopped.is_set() + assert not caplog.records + + +@pytest.mark.asyncio +@pytest.mark.parametrize("manager_factory", MANAGER_FACTORIES) +async def test_stop_before_start_and_repeated_stop_are_safe(manager_factory): + client = IdleClient() + manager = manager_factory(client) + + await asyncio.wait_for(manager.stop(), timeout=0.1) + await asyncio.wait_for(manager.stop(), timeout=0.1) + + assert client.closed is True + assert not client.messages_started.is_set() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("manager_factory", MANAGER_FACTORIES) +async def test_start_while_running_still_fails(manager_factory): + client = IdleClient() + manager = manager_factory(client) + + await manager.start() + await asyncio.wait_for(client.messages_started.wait(), timeout=0.1) + try: + with pytest.raises(RuntimeError, match="already running"): + await manager.start() + finally: + await asyncio.wait_for(manager.stop(), timeout=0.1) + + assert client.closed is True + assert client.messages_stopped.is_set()