Fix idle ConnectionManager shutdown

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

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:

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.

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

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`.

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:

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:

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:

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()