From 018259dd0a7cb91b9d38b29e19a2556dd0bfbc08 Mon Sep 17 00:00:00 2001 From: Michael Brown Date: Fri, 1 May 2026 04:39:51 +0800 Subject: [PATCH 1/2] Allow subscriptions to all subaccount dashboards --- vanta_api/websocket_server.py | 86 +++++++++++++++++------------------ 1 file changed, 43 insertions(+), 43 deletions(-) diff --git a/vanta_api/websocket_server.py b/vanta_api/websocket_server.py index bbb811740..cc5dcc0ff 100644 --- a/vanta_api/websocket_server.py +++ b/vanta_api/websocket_server.py @@ -99,6 +99,7 @@ class WebSocketServerClient: tier: int = 0 sequence_number: int = 0 subscribe_broadcasts: bool = False + subscribe_all_dashboard_updates: bool = False dashboard_subscriptions: dict[str, DashboardSubscription] = field(default_factory=dict) async def send(self, message:dict) -> None: @@ -435,6 +436,11 @@ def _process_dashboard_update_queue(self) -> None: synthetic_hotkey = self._dashboard_update_queue.get() for client in list(self._clients.values()): subscription = client.dashboard_subscriptions.get(synthetic_hotkey) + + if (subscription is None) and client.subscribe_all_dashboard_updates: + subscription = DashboardSubscription() + client.dashboard_subscriptions[synthetic_hotkey] = subscription + if subscription is not None: self._thread_pool.submit(self._send_dashboard_update, synthetic_hotkey, client, subscription) except Exception as e: @@ -585,47 +591,11 @@ async def _handle_client(self, websocket) -> None: bt.logging.info(f"WebSocketServer: Client {client_id} authenticated successfully with tier {api_key_tier}") - # For entity clients (tier >= 200), auto-subscribe to all active subaccounts - hl_mappings = {} - subscribed_subaccounts = 0 - - if api_key_tier >= ValiConfig.SUBACCOUNT_SUBSCRIPTION_TIER: - entity_hotkey = self.api_key_to_alias.get(api_key) - if entity_hotkey: - try: - loop = asyncio.get_running_loop() - entity_data = await loop.run_in_executor( - self._thread_pool, - self._entity_client.get_entity_data, - entity_hotkey - ) - if entity_data: - subaccounts = entity_data.get('subaccounts', {}) - for sub_id, sub_info in subaccounts.items(): - if sub_info.get('status') not in ('active', 'admin'): - continue - synthetic_hotkey = sub_info.get('synthetic_hotkey') or f"{entity_hotkey}_{sub_id}" - client.dashboard_subscriptions[synthetic_hotkey] = DashboardSubscription() - subscribed_subaccounts += 1 - hl_addr = sub_info.get('hl_address') - if hl_addr: - hl_mappings[hl_addr] = synthetic_hotkey - if subscribed_subaccounts > 0: - bt.logging.info( - f"WebSocketServer: Auto-subscribed client {client_id} to " - f"{subscribed_subaccounts} subaccounts for entity {entity_hotkey}" - ) - except Exception as e: - bt.logging.error(f"WebSocketServer: Error auto-subscribing entity {entity_hotkey}: {e}") - bt.logging.error(traceback.format_exc()) - await websocket.send(json.dumps({ "status": "success", "message": "Authentication successful.", "current_sequence": 0, "tier": api_key_tier, - "subscribed_subaccounts": subscribed_subaccounts, - "hl_mappings": hl_mappings })) # Process client messages (subscriptions, etc.) @@ -654,7 +624,7 @@ async def _handle_client(self, websocket) -> None: "type": "subscription_status", "status": "success", "all": True, - "action": "subscribe", + "action": message_type, "sender_timestamp": data.get("sender_timestamp", 0), "received_timestamp": TimeUtil.now_in_millis() })) @@ -668,7 +638,7 @@ async def _handle_client(self, websocket) -> None: "type": "subscription_status", "status": "success", "all": True, - "action": "unsubscribe" + "action": message_type })) elif message_type == "subscribe_subaccount": @@ -682,7 +652,7 @@ async def _handle_client(self, websocket) -> None: await websocket.send(json.dumps({ "type": "subscription_status", "status": "error", - "action": "subscribe_subaccount", + "action": message_type, "message": "Missing synthetic_hotkey parameter" })) else: @@ -705,14 +675,14 @@ async def _handle_client(self, websocket) -> None: await websocket.send(json.dumps({ "type": "subscription_status", "status": "success", - "action": "subscribe_subaccount", + "action": message_type, "subscribed_to": synthetic_hotkey })) elif client.tier < ValiConfig.SUBACCOUNT_SUBSCRIPTION_TIER: await websocket.send(json.dumps({ "type": "subscription_status", "status": "error", - "action": "subscribe_subaccount", + "action": message_type, "synthetic_hotkey": synthetic_hotkey, "message": "Subaccount subscriptions require tier 200 access.", "code": "INSUFFICIENT_TIER" @@ -728,7 +698,7 @@ async def _handle_client(self, websocket) -> None: await websocket.send(json.dumps({ "type": "subscription_status", "status": "success", - "action": "subscribe_subaccount", + "action": message_type, "subscribed_to": synthetic_hotkey })) @@ -739,10 +709,40 @@ async def _handle_client(self, websocket) -> None: await websocket.send(json.dumps({ "type": "subscription_status", "status": "success", - "action": "unsubscribe_subaccount", + "action": message_type, "synthetic_hotkey": synthetic_hotkey })) + elif message_type == "subscribe_subaccounts": + if client.tier < ValiConfig.SUBACCOUNT_SUBSCRIPTION_TIER: + await websocket.send(json.dumps({ + "type": "subscription_status", + "status": "error", + "action": message_type, + "message": "Subaccount subscriptions require tier 200 access.", + "code": "INSUFFICIENT_TIER" + })) + else: + client.subscribe_all_dashboard_updates = True + bt.logging.info( + f"WebSocketServer: Client {client_id} subscribed to all subaccounts") + await websocket.send(json.dumps({ + "type": "subscription_status", + "status": "success", + "action": message_type, + })) + + elif message_type == "unsubscribe_subaccounts": + client.subscribe_all_dashboard_updates = False + client.dashboard_subscriptions.clear() + bt.logging.info( + f"WebSocketServer: Client {client_id} unsubscribed from all subaccounts") + await websocket.send(json.dumps({ + "type": "subscription_status", + "status": "success", + "action": message_type, + })) + except websockets.exceptions.ConnectionClosed: break except json.JSONDecodeError: From ceaae97a4564faea36cb2d7e1eb58c8b5a3b91db Mon Sep 17 00:00:00 2001 From: Michael Brown Date: Wed, 13 May 2026 16:18:21 +0800 Subject: [PATCH 2/2] Ensure client has permission for synthetic hotkey subscriptions --- vanta_api/websocket_server.py | 27 ++++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/vanta_api/websocket_server.py b/vanta_api/websocket_server.py index cc5dcc0ff..b6c0bd2cf 100644 --- a/vanta_api/websocket_server.py +++ b/vanta_api/websocket_server.py @@ -96,12 +96,16 @@ class WebSocketServerClient: client_id: int websocket: WebSocketServerProtocol api_key: str + entity_hotkey: str tier: int = 0 sequence_number: int = 0 subscribe_broadcasts: bool = False subscribe_all_dashboard_updates: bool = False dashboard_subscriptions: dict[str, DashboardSubscription] = field(default_factory=dict) + def hotkey_matches(self, hotkey: str) -> bool: + return self.entity_hotkey and hotkey.startswith(self.entity_hotkey) + async def send(self, message:dict) -> None: serialized_message = json.dumps( { @@ -437,7 +441,9 @@ def _process_dashboard_update_queue(self) -> None: for client in list(self._clients.values()): subscription = client.dashboard_subscriptions.get(synthetic_hotkey) - if (subscription is None) and client.subscribe_all_dashboard_updates: + if ((subscription is None) and + client.subscribe_all_dashboard_updates and + client.hotkey_matches(synthetic_hotkey)): subscription = DashboardSubscription() client.dashboard_subscriptions[synthetic_hotkey] = subscription @@ -585,7 +591,14 @@ async def _handle_client(self, websocket) -> None: f"{api_key_alias} to make room for new client {client_id}") # Register client - client = WebSocketServerClient(client_id=client_id, websocket=websocket, api_key=api_key, tier=api_key_tier) + entity_hotkey = self.api_key_to_alias.get(api_key) + client = WebSocketServerClient( + client_id=client_id, + websocket=websocket, + api_key=api_key, + tier=api_key_tier, + entity_hotkey=entity_hotkey + ) self._clients[client_id] = client self._api_key_client_ids[api_key].append(client_id) @@ -687,7 +700,7 @@ async def _handle_client(self, websocket) -> None: "message": "Subaccount subscriptions require tier 200 access.", "code": "INSUFFICIENT_TIER" })) - else: + elif client.hotkey_matches(synthetic_hotkey): client.dashboard_subscriptions[synthetic_hotkey] = DashboardSubscription( positions_time_ms=positions_time_ms, limit_orders_time_ms=limit_orders_time_ms, @@ -701,6 +714,14 @@ async def _handle_client(self, websocket) -> None: "action": message_type, "subscribed_to": synthetic_hotkey })) + else: + await websocket.send(json.dumps({ + "type": "subscription_status", + "status": "error", + "action": message_type, + "synthetic_hotkey": synthetic_hotkey, + "message": "Subaccount not authorized for this API key.", + })) elif message_type == "unsubscribe_subaccount": synthetic_hotkey = data.get("synthetic_hotkey")