diff --git a/examples/websocket/ws_auth.py b/examples/websocket/ws_auth.py new file mode 100644 index 0000000..e571cbf --- /dev/null +++ b/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()) diff --git a/lighter/ws_client.py b/lighter/ws_client.py index 0638875..f783948 100644 --- a/lighter/ws_client.py +++ b/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: