Add auth_token to WsClient, sent as Authorization header on connect

examples/websocket/ws_auth.py

@@ -0,0 +1,49 @@
+import json
+import logging
+import asyncio
+import lighter
+from examples.utils import get_api_key_config
+
+
+def on_order_book_update(market_id, order_book):
+    logging.info(f"Order book {market_id}:\n{json.dumps(order_book, indent=2)}")
+
+
+def on_account_update(account_id, account):
+    logging.info(f"Account {account_id}:\n{json.dumps(account, indent=2)}")
+
+
+async def main():
+    logging.basicConfig(level=logging.INFO)
+
+    profile, account_index, private_keys = get_api_key_config()
+    client = lighter.SignerClient(
+        url=profile.api_url,
+        account_index=account_index,
+        api_private_keys=private_keys,
+        chain_id=profile.chain_id,
+    )
+
+    # Authenticating the connection applies your account tier's websocket limits instead of the public ones.
+    # The token only has to be valid when the connection opens.
+    auth_token, err = client.create_auth_token_with_expiry()
+    if err is not None:
+        raise Exception(err)
+
+    ws_client = lighter.WsClient(
+        ws_url=profile.ws_url,
+        order_book_ids=[0, 1],
+        account_ids=[account_index],
+        on_order_book_update=on_order_book_update,
+        on_account_update=on_account_update,
+        auth_token=auth_token,
+    )
+
+    try:
+        await ws_client.run_async()
+    finally:
+        await client.close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

lighter/ws_client.py

@@ -18,6 +18,7 @@ def __init__(
         ws_url=None,
         on_shutdown=None,
         on_unhandled_message=None,
+        auth_token=None,
     ):
         if ws_url is not None:
             self.base_url = ws_url.rstrip("/")
@@ -41,6 +42,7 @@ def __init__(
         self.on_account_update = on_account_update
         self.on_shutdown = on_shutdown
         self.on_unhandled_message = on_unhandled_message
+        self.auth_token = auth_token
 
         self.ws = None
 
@@ -175,15 +177,18 @@ def on_error(self, ws, error):
     def on_close(self, ws, close_status_code, close_msg):
         raise Exception(f"Closed: {close_status_code} {close_msg}")
 
+    def _auth_headers(self):
+        return {"Authorization": self.auth_token} if self.auth_token else None
+
     def run(self):
-        ws = connect(self.base_url)
+        ws = connect(self.base_url, additional_headers=self._auth_headers())
         self.ws = ws
 
         for message in ws:
             self.on_message(ws, message)
 
     async def run_async(self):
-        ws = await connect_async(self.base_url)
+        ws = await connect_async(self.base_url, extra_headers=self._auth_headers())
         self.ws = ws
 
         async for message in ws: