fix: preserve WebSocket error subscription context

README.md

@@ -194,6 +194,9 @@ except APIError as e:
 separately with `SUBSCRIPTION_TYPE_ORDER_SNAPSHOT` and a distinct request ID. A
 successful snapshot ends with an `eof: true` frame; failures use the `error` handler.
 
+For server error frames, both streams emit a `WebSocketError` with optional
+`request_id` and `subscription_type` attributes when supplied by the server.
+
 ```python
 import asyncio
 import os

polymarket_us/errors.py

@@ -123,9 +123,16 @@ class InternalServerError(APIStatusError):
 class WebSocketError(PolymarketUSError):
     """WebSocket-related error."""
 
-    def __init__(self, message: str, request_id: str | None = None) -> None:
+    def __init__(
+        self,
+        message: str,
+        request_id: str | None = None,
+        *,
+        subscription_type: str | None = None,
+    ) -> None:
         super().__init__(message)
         self.request_id = request_id
+        self.subscription_type = subscription_type
 
     def __repr__(self) -> str:
         return f"WebSocketError(message={str(self)!r}, request_id={self.request_id!r})"

polymarket_us/websocket/markets.py

@@ -72,7 +72,11 @@ def _handle_message(self, data: str) -> None:
         elif "error" in message:
             self._emit(
                 "error",
-                WebSocketError(message["error"], message.get("requestId")),
+                WebSocketError(
+                    message["error"],
+                    message.get("requestId"),
+                    subscription_type=message.get("subscriptionType"),
+                ),
             )
         elif "marketData" in message:
             self._emit("market_data", message)

polymarket_us/websocket/private.py

@@ -79,7 +79,11 @@ def _handle_message(self, data: str) -> None:
         elif message.get("error"):
             self._emit(
                 "error",
-                WebSocketError(message["error"], message.get("requestId")),
+                WebSocketError(
+                    message["error"],
+                    message.get("requestId"),
+                    subscription_type=message.get("subscriptionType"),
+                ),
             )
         elif "orderSubscriptionSnapshot" in message or "ordersSnapshot" in message:
             self._emit("order_snapshot", message)

polymarket_us/websocket/types.py

@@ -272,6 +272,7 @@ class WebSocketErrorMessage(TypedDict, total=False):
 
     requestId: str
     error: str
+    subscriptionType: str
 
 
 PrivateMessage = (

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "polymarket-us"
-version = "2.4.0"
+version = "2.4.1"
 description = "Polymarket US Python SDK"
 readme = "README.md"
 license = "MIT"

tests/test_errors.py

@@ -178,3 +178,16 @@ def test_stores_request_id(self) -> None:
         """Should store request_id."""
         error = WebSocketError("test", request_id="req-123")
         assert error.request_id == "req-123"
+        assert error.subscription_type is None
+
+    def test_subscription_context_preserves_existing_error_behavior(self) -> None:
+        error = WebSocketError(
+            "subscription failed", "req-123", subscription_type="SUBSCRIPTION_TYPE_ORDER"
+        )
+        assert error.subscription_type == "SUBSCRIPTION_TYPE_ORDER"
+        assert error.request_id == "req-123"
+        assert error.args == ("subscription failed",)
+        assert str(error) == "subscription failed"
+        assert repr(error) == (
+            "WebSocketError(message='subscription failed', request_id='req-123')"
+        )

tests/test_websocket.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import base64
+import json
 from collections.abc import AsyncIterator
 from typing import Literal
 
@@ -11,9 +12,10 @@
 from websockets.http11 import Request
 
 from polymarket_us import AuthenticationError, PolymarketUS
-from polymarket_us.errors import PolymarketUSError
+from polymarket_us.errors import PolymarketUSError, WebSocketError
 from polymarket_us.websocket import MarketsWebSocket, PrivateWebSocket
 from polymarket_us.websocket.base import BaseWebSocket
+from polymarket_us.websocket.types import WebSocketErrorMessage
 
 TEST_SECRET_KEY = "nWGxne/9WmC6hEr0kuwsxERJxWl7MmkZcDusAxyuf2A="
 
@@ -84,6 +86,50 @@ async def handler(peer: ServerConnection) -> None:
                 finally:
                     await asyncio.wait_for(ws.close(), timeout=5)
 
+    @pytest.mark.parametrize(
+        "subscription_type",
+        ["known", "SUBSCRIPTION_TYPE_UNSPECIFIED", "SUBSCRIPTION_TYPE_FUTURE", None],
+    )
+    @pytest.mark.parametrize("request_id", ["req-123", None])
+    async def test_error_callback_preserves_subscription_context(
+        self,
+        connection: tuple[BaseWebSocket, ServerConnection],
+        stream: Literal["markets", "private"],
+        subscription_type: str | None,
+        request_id: str | None,
+    ) -> None:
+        ws, peer = connection
+        if subscription_type == "known":
+            subscription_type = (
+                "SUBSCRIPTION_TYPE_MARKET_DATA"
+                if stream == "markets"
+                else "SUBSCRIPTION_TYPE_ORDER"
+            )
+        frame: WebSocketErrorMessage = {"error": "subscription failed"}
+        if subscription_type is not None:
+            frame["subscriptionType"] = subscription_type
+        if request_id is not None:
+            frame["requestId"] = request_id
+        errors: list[PolymarketUSError] = []
+        messages: list[object] = []
+        heartbeat = asyncio.Event()
+        ws.on("error", errors.append)
+        ws.on("message", messages.append)
+        ws.on("heartbeat", heartbeat.set)
+
+        await asyncio.wait_for(peer.send(json.dumps(frame)), timeout=5)
+        await asyncio.wait_for(peer.send('{"heartbeat": {}}'), timeout=5)
+        await asyncio.wait_for(heartbeat.wait(), timeout=5)
+
+        assert messages == [frame, {"heartbeat": {}}]
+        assert len(errors) == 1
+        error = errors[0]
+        assert isinstance(error, WebSocketError)
+        assert str(error) == "subscription failed"
+        assert error.request_id == request_id
+        assert error.subscription_type == subscription_type
+        assert ws.is_connected
+
     @pytest.mark.parametrize("close_code", [1000, 1001, 1011])
     async def test_remote_close_emits_once(
         self, connection: tuple[BaseWebSocket, ServerConnection], close_code: int

tests/types/websocket_errors.py

@@ -0,0 +1,17 @@
+"""Typed consumers of WebSocket error context, including unknown subscriptions."""
+
+from polymarket_us import WebSocketError
+from polymarket_us.websocket import WebSocketErrorMessage
+
+legacy_frame: WebSocketErrorMessage = {"error": "subscription failed"}
+frame: WebSocketErrorMessage = {
+    "error": "subscription failed",
+    "requestId": "req-123",
+    "subscriptionType": "SUBSCRIPTION_TYPE_FUTURE",
+}
+error = WebSocketError(
+    frame["error"],
+    frame.get("requestId"),
+    subscription_type=frame.get("subscriptionType"),
+)
+subscription_type: str | None = error.subscription_type