WebsocketManager lacks reconnection and subscription restoration

[USEONGEE]

The current WebsocketManager implementation has the following limitations:

1. No Automatic Reconnection:

• There is no logic to automatically reconnect when the connection is closed (on_close) or an error occurs (on_error).
• In case of network instability or server restarts, the client must manually implement reconnection.

2. No Subscription Restoration:

• Even if the connection is re-established, existing active_subscriptions / queued_subscriptions are not automatically re-subscribed.

[mifth]

Hey, @USEONGEE . I've found that there is a relatively new kwarg reconnect in ws.run_forever(). It's None by default. I set it to 5 and my conncetion is working during 1 day. Could you please test it?

import hyperliquid
import hyperliquid.info
from hyperliquid.utils import constants

class MyWebsocketManager(hyperliquid.info.WebsocketManager):
    def __init__(self, base_url: str, *args, **kwargs):
        super().__init__(base_url, *args, **kwargs)

    def run(self):
        self.ping_sender.start()
        self.ws.run_forever(reconnect=5)  # Set reconnect to 5 seconds


def _replace_ws_manager(info):
    try:
        if info.ws_manager:
            info.ws_manager.stop()
            info.ws_manager.join()
    except Exception as e:
        logging.error("Error stopping old websocket manager: %s", e)
        logging.error(traceback.format_exc())

    info.ws_manager = MyWebsocketManager(base_url=constants.MAINNET_API_URL)
    info.ws_manager.start()


# Test
 info: hyperliquid.info.Info = hyperliquid.info.Info(base_url=constants.MAINNET_API_URL)
_replace_ws_manager(info)

[mifth]

I've created a new pull request as it's pretty stable. #255

[USEONGEE]

I've created a new pull request as it's pretty stable. #255

Good men. I think your solution is better.

[mifth]

This variant should be perfect. We also need to re-subscribe all subscriptions again during the reconnection.

import hyperliquid
import hyperliquid.info
from hyperliquid.utils import constants
import json
from typing import Optional


# Fixed websocket manager to allow restart on stop
class MyWebsocketManager(hyperliquid.info.WebsocketManager):
    def __init__(self, base_url: str, *args, **kwargs):
        super().__init__(base_url, *args, **kwargs)

        # Subscription map. We need it for reconnection.
        self.subscription_map: dict[int, hyperliquid.info.Subscription] = {}  # Key - subscription ID. Value - Subscription.

        # Reconnection event handler
        if self.ws and self.ws.on_reconnect is None:
            self.ws.on_reconnect = self.on_reconnect
        else:
            logging.error("Websocket is not initialized properly in MyWebsocketManager. Cannot set on_reconnect handler!!!")

    # Override the run method
    def run(self):
        self.ping_sender.start()
        self.ws.run_forever(reconnect=5)


    def on_reconnect(self, _ws):
        # Unsubscribe and resubscribe all active subscriptions on reconnect
        subscriptions_copy: list[hyperliquid.info.Subscription] = list(self.subscription_map.values())
        for subscription in subscriptions_copy:
            try:
                self.ws.send(json.dumps({"method": "unsubscribe", "subscription": subscription}))
                logging.info(f"Unsubscribed from {subscription} on reconnect.")
            except Exception as e:
                logging.error(f"Error unsubscribing from {subscription} on reconnect: {e}")
                logging.error(traceback.format_exc())

            try:
                self.ws.send(json.dumps({"method": "subscribe", "subscription": subscription}))
                logging.info(f"Resubscribed to {subscription} on reconnect.")
            except Exception as e:
                logging.error(f"Error resubscribing to {subscription} on reconnect: {e}")
                logging.error(traceback.format_exc())

    
    # Override the subscribe method
    def subscribe(self, subscription: hyperliquid.info.Subscription, callback, 
                  subscription_id: Optional[int] = None) -> int:
        subscription_id: int = super().subscribe(subscription, callback, subscription_id)
        self.subscription_map[subscription_id] = subscription
        
        return subscription_id


    # Override the unsubscribe method
    def unsubscribe(self, subscription: hyperliquid.info.Subscription, subscription_id: int) -> bool:
        is_unsubscribed: bool =  super().unsubscribe(subscription, subscription_id)
        if is_unsubscribed and subscription_id in self.subscription_map:
            del self.subscription_map[subscription_id]

        return is_unsubscribed


    @staticmethod
    def replace_ws_manager(info: hyperliquid.info.Info):
        try:
            if info.ws_manager:
                info.ws_manager.stop()
        except Exception as e:
            logging.error("Error stopping old websocket manager on stop() : %s", e)
            logging.error(traceback.format_exc())

        try:
            if info.ws_manager:
                info.ws_manager.join()
        except Exception as e:
            logging.error("Error stopping old websocket manager on join(): %s", e)
            logging.error(traceback.format_exc())

        info.ws_manager = MyWebsocketManager(base_url=constants.MAINNET_API_URL)
        info.ws_manager.start()


# Test
info: hyperliquid.info.Info = hyperliquid.info.Info(base_url=constants.MAINNET_API_URL)
MyWebsocketManager.replace_ws_manager(info)