Check single-subscription channels before queueing, not during replay

hyperliquid/websocket_manager.py

@@ -127,7 +127,11 @@ def on_message(self, _ws, message):
     def on_open(self, _ws):
         logging.debug("on_open")
         self.ws_ready = True
-        for subscription, active_subscription in self.queued_subscriptions:
+        # Drain the queue before replaying it: subscribe() consults it for the
+        # single-subscription channels, and leaving entries behind would also
+        # replay them again on a later on_open.
+        queued_subscriptions, self.queued_subscriptions = self.queued_subscriptions, []
+        for subscription, active_subscription in queued_subscriptions:
             self.subscribe(subscription, active_subscription.callback, active_subscription.subscription_id)
 
     def subscribe(
@@ -136,16 +140,19 @@ def subscribe(
         if subscription_id is None:
             self.subscription_id_counter += 1
             subscription_id = self.subscription_id_counter
+        identifier = subscription_to_identifier(subscription)
+        if identifier == "userEvents" or identifier == "orderUpdates":
+            # TODO: ideally the userEvent and orderUpdates messages would include the user so that we can multiplex
+            # Queued subscriptions count too, otherwise the duplicate is only caught while on_open replays the
+            # queue, where it aborts the replay and silently drops every subscription behind it.
+            already_queued = any(subscription_to_identifier(s) == identifier for s, _ in self.queued_subscriptions)
+            if len(self.active_subscriptions[identifier]) != 0 or already_queued:
+                raise NotImplementedError(f"Cannot subscribe to {identifier} multiple times")
         if not self.ws_ready:
             logging.debug("enqueueing subscription")
             self.queued_subscriptions.append((subscription, ActiveSubscription(callback, subscription_id)))
         else:
             logging.debug("subscribing")
-            identifier = subscription_to_identifier(subscription)
-            if identifier == "userEvents" or identifier == "orderUpdates":
-                # TODO: ideally the userEvent and orderUpdates messages would include the user so that we can multiplex
-                if len(self.active_subscriptions[identifier]) != 0:
-                    raise NotImplementedError(f"Cannot subscribe to {identifier} multiple times")
             self.active_subscriptions[identifier].append(ActiveSubscription(callback, subscription_id))
             self.ws.send(json.dumps({"method": "subscribe", "subscription": subscription}))
         return subscription_id

tests/websocket_manager_test.py

@@ -0,0 +1,77 @@
+import json
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+
+from hyperliquid.websocket_manager import WebsocketManager
+
+
+def make_manager(ws_ready: bool):
+    """A WebsocketManager with a stub socket, so no connection is opened."""
+    ws_manager = WebsocketManager.__new__(WebsocketManager)
+    ws_manager.subscription_id_counter = 0
+    ws_manager.ws_ready = ws_ready
+    ws_manager.queued_subscriptions = []
+    ws_manager.active_subscriptions = defaultdict(list)
+    sent = []
+    ws_manager.ws = SimpleNamespace(send=sent.append)
+    return ws_manager, sent
+
+
+def callback(_msg):
+    pass
+
+
+def test_duplicate_single_subscription_raises_while_queued():
+    # userEvents and orderUpdates cannot be multiplexed. Before this was checked on
+    # the queued path the duplicate was only caught later, inside on_open.
+    ws_manager, _ = make_manager(ws_ready=False)
+    ws_manager.subscribe({"type": "userEvents"}, callback)
+
+    with pytest.raises(NotImplementedError):
+        ws_manager.subscribe({"type": "userEvents"}, callback)
+
+
+def test_duplicate_single_subscription_raises_when_connected():
+    ws_manager, _ = make_manager(ws_ready=True)
+    ws_manager.subscribe({"type": "orderUpdates"}, callback)
+
+    with pytest.raises(NotImplementedError):
+        ws_manager.subscribe({"type": "orderUpdates"}, callback)
+
+
+def test_rejected_duplicate_does_not_drop_later_subscriptions():
+    # The duplicate used to surface inside on_open, which aborted the replay and
+    # silently dropped every subscription queued behind it.
+    ws_manager, sent = make_manager(ws_ready=False)
+    ws_manager.subscribe({"type": "userEvents"}, callback)
+    with pytest.raises(NotImplementedError):
+        ws_manager.subscribe({"type": "userEvents"}, callback)
+    ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback)
+
+    ws_manager.on_open(None)
+
+    assert len(ws_manager.active_subscriptions["l2Book:eth"]) == 1
+    assert [json.loads(msg)["subscription"]["type"] for msg in sent] == ["userEvents", "l2Book"]
+
+
+def test_on_open_replays_and_clears_the_queue():
+    ws_manager, sent = make_manager(ws_ready=False)
+    ws_manager.subscribe({"type": "l2Book", "coin": "BTC"}, callback)
+    ws_manager.subscribe({"type": "trades", "coin": "ETH"}, callback)
+
+    ws_manager.on_open(None)
+
+    assert len(sent) == 2
+    assert len(ws_manager.active_subscriptions["l2Book:btc"]) == 1
+    assert len(ws_manager.active_subscriptions["trades:eth"]) == 1
+    assert ws_manager.queued_subscriptions == []
+
+
+def test_multiplexable_channel_still_accepts_several_callbacks():
+    ws_manager, _ = make_manager(ws_ready=True)
+    ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback)
+    ws_manager.subscribe({"type": "l2Book", "coin": "ETH"}, callback)
+
+    assert len(ws_manager.active_subscriptions["l2Book:eth"]) == 2