""" Connection management functionality for real-time client. Author: @TexasCoding Date: 2025-08-02 Overview: Provides connection management functionality for the ProjectX real-time client, including SignalR hub setup, connection establishment, reconnection handling, and secure JWT token authentication. Key Features: - Dual-hub SignalR connection setup and management - Automatic reconnection with exponential backoff - JWT token authentication via URL query parameter (ProjectX Gateway requirement) - Connection health monitoring and error handling - Thread-safe operations with proper lock management - Comprehensive connection statistics and health tracking Connection Management Capabilities: - SignalR hub setup with ProjectX Gateway configuration - Connection establishment and health monitoring - Automatic reconnection with configurable intervals - JWT token authentication via URL query parameter (ProjectX Gateway requirement) - JWT token refresh and reconnection handling - Connection event handling and error processing - Statistics tracking and health reporting Example Usage: The functionality of this mixin is consumed through a `ProjectXRealtimeClient` instance. For most use cases, this is handled automatically by the `TradingSuite`. ```python # The following demonstrates the low-level connection lifecycle managed by this mixin. # Note: In a typical application, you would use TradingSuite, which handles this. from project_x_py import create_realtime_client # 1. Initialization (handled by factory or TradingSuite) # realtime_client = await create_realtime_client(jwt, account_id) # 2. Connection (handled by TradingSuite.create() or client.connect()) # if await realtime_client.connect(): # print("Successfully connected to both User and Market hubs.") # # 3. Health Monitoring # if realtime_client.is_connected(): # stats = realtime_client.get_stats() # print(f"Events received so far: {stats['events_received']}") # # 4. Token Refresh (if managing tokens manually) # # new_token = await get_new_token() # # await realtime_client.update_jwt_token(new_token) # # 5. Disconnection (handled by TradingSuite context manager or client.disconnect()) # await realtime_client.disconnect() ``` See Also: - `realtime.core.ProjectXRealtimeClient` - `realtime.event_handling.EventHandlingMixin` - `realtime.subscriptions.SubscriptionsMixin` """ import asyncio from datetime import datetime from typing import TYPE_CHECKING, Any from project_x_py.utils import ( LogContext, LogMessages, ProjectXLogger, handle_errors, ) try: from signalrcore.hub_connection_builder import HubConnectionBuilder except ImportError: HubConnectionBuilder = None if TYPE_CHECKING: from project_x_py.types import ProjectXRealtimeClientProtocol logger = ProjectXLogger.get_logger(__name__) class ConnectionManagementMixin: """Mixin for connection management functionality.""" def __init__(self) -> None: """Initialize connection management attributes.""" super().__init__() self._loop: asyncio.AbstractEventLoop | None = None @handle_errors("setup connections") async def setup_connections(self: "ProjectXRealtimeClientProtocol") -> None: """ Set up SignalR hub connections with ProjectX Gateway configuration. Initializes both user and market hub connections with secure JWT authentication, proper event handlers, automatic reconnection, and ProjectX-specific event mappings. Must be called before connect() or is called automatically on first connect(). Hub Configuration: - User Hub: Account, position, order, and trade events - Market Hub: Quote, trade, and market depth events - Both hubs: Automatic reconnection with exponential backoff - Keep-alive: 10 second interval - Reconnect intervals: [1, 3, 5, 5, 5, 5] seconds Authentication: - JWT token in URL query parameter (ProjectX Gateway requirement) Event Mappings: User Hub Events: - GatewayUserAccount -> account_update - GatewayUserPosition -> position_update - GatewayUserOrder -> order_update - GatewayUserTrade -> trade_execution Market Hub Events: - GatewayQuote -> quote_update - GatewayTrade -> market_trade - GatewayDepth -> market_depth Raises: ImportError: If signalrcore package is not installed Exception: If connection setup fails Note: This method is idempotent - safe to call multiple times. Sets self.setup_complete = True when successful. """ with LogContext( logger, operation="setup_connections", user_hub=self.user_hub_url, market_hub=self.market_hub_url, ): logger.debug(LogMessages.WS_CONNECT, extra={"phase": "setup"}) if HubConnectionBuilder is None: raise ImportError("signalrcore is required for real-time functionality") async with self._connection_lock: logger.info( "Using URL query parameter for JWT authentication (ProjectX Gateway requirement)" ) # Build user hub connection with JWT as query parameter # ProjectX Gateway requires auth tokens in URL for WebSocket connections user_url_with_token = ( f"{self.user_hub_url}?access_token={self.jwt_token}" ) self.user_connection = ( HubConnectionBuilder() .with_url(user_url_with_token) .configure_logging( logger.level, socket_trace=False, handler=None, ) .with_automatic_reconnect( { "type": "interval", "keep_alive_interval": 10, "intervals": [1, 3, 5, 5, 5, 5], } ) .build() ) # Build market hub connection with JWT as query parameter market_url_with_token = ( f"{self.market_hub_url}?access_token={self.jwt_token}" ) self.market_connection = ( HubConnectionBuilder() .with_url(market_url_with_token) .configure_logging( logger.level, socket_trace=False, handler=None, # Use None to avoid duplicate logging ) .with_automatic_reconnect( { "type": "interval", "keep_alive_interval": 10, "intervals": [1, 3, 5, 5, 5, 5], } ) .build() ) # Set up connection event handlers assert self.user_connection is not None assert self.market_connection is not None self.user_connection.on_open(lambda: self._on_user_hub_open()) self.user_connection.on_close(lambda: self._on_user_hub_close()) self.user_connection.on_error( lambda data: self._on_connection_error("user", data) ) self.market_connection.on_open(lambda: self._on_market_hub_open()) self.market_connection.on_close(lambda: self._on_market_hub_close()) self.market_connection.on_error( lambda data: self._on_connection_error("market", data) ) # Set up ProjectX Gateway event handlers (per official documentation) # User Hub Events self.user_connection.on( "GatewayUserAccount", self._forward_account_update ) self.user_connection.on( "GatewayUserPosition", self._forward_position_update ) self.user_connection.on("GatewayUserOrder", self._forward_order_update) self.user_connection.on( "GatewayUserTrade", self._forward_trade_execution ) # Market Hub Events self.market_connection.on("GatewayQuote", self._forward_quote_update) self.market_connection.on("GatewayTrade", self._forward_market_trade) self.market_connection.on("GatewayDepth", self._forward_market_depth) logger.debug( LogMessages.WS_CONNECTED, extra={"phase": "setup_complete"} ) self.setup_complete = True @handle_errors("connect", reraise=False, default_return=False) async def connect(self: "ProjectXRealtimeClientProtocol") -> bool: """ Connect to ProjectX Gateway SignalR hubs asynchronously. Establishes connections to both user and market hubs, enabling real-time event streaming. Connections are made concurrently for efficiency. Returns: bool: True if both hubs connected successfully, False otherwise Connection Process: 1. Sets up connections if not already done 2. Stores event loop for cross-thread operations 3. Starts user hub connection 4. Starts market hub connection 5. Waits for connection establishment 6. Updates connection statistics Example: >>> # V3.1: TradingSuite handles connection automatically >>> suite = await TradingSuite.create("MNQ", timeframes=["1min"]) >>> print(f"Connected: {suite.realtime_client.is_connected()}") >>> # V3.1: Direct usage (advanced) >>> # client = ProjectXRealtimeClient(jwt_token, account_id) >>> # if await client.connect(): >>> # await client.subscribe_market_data(["MNQ", "ES"]) Side Effects: - Sets self.user_connected and self.market_connected flags - Updates connection statistics - Stores event loop reference Note: - Both hubs must connect for success - SignalR connections run in thread executor for async compatibility - Automatic reconnection is configured but initial connect may fail """ with LogContext( logger, operation="connect", account_id=self.account_id, ): if not self.setup_complete: await self.setup_connections() # Store the event loop for cross-thread task scheduling try: self._loop = asyncio.get_running_loop() except RuntimeError: logger.error("No running event loop found.") return False logger.debug(LogMessages.WS_CONNECT) async with self._connection_lock: # Start both connections if self.user_connection: await self._start_connection_async(self.user_connection, "user") else: logger.error( LogMessages.WS_ERROR, extra={"error": "User connection not available"}, ) return False if self.market_connection: await self._start_connection_async(self.market_connection, "market") else: logger.error( LogMessages.WS_ERROR, extra={"error": "Market connection not available"}, ) return False # Wait for connections to establish try: await asyncio.wait_for( asyncio.gather( self.user_hub_ready.wait(), self.market_hub_ready.wait() ), timeout=10.0, ) except TimeoutError: logger.error( LogMessages.WS_ERROR, extra={ "error": "Connection attempt timed out after 10 seconds." }, ) return False if self.user_connected and self.market_connected: self.stats["connected_time"] = datetime.now() logger.debug(LogMessages.WS_CONNECTED) return True else: logger.error( LogMessages.WS_ERROR, extra={"error": "Failed to establish all connections"}, ) return False @handle_errors("start connection") async def _start_connection_async( self: "ProjectXRealtimeClientProtocol", connection: Any, name: str ) -> None: """ Start a SignalR connection asynchronously. Wraps the synchronous SignalR start() method to work with asyncio by running it in a thread executor. Args: connection: SignalR HubConnection instance to start name (str): Hub name for logging ("user" or "market") Note: This is an internal method that bridges sync SignalR with async code. """ # SignalR connections are synchronous, so we run them in executor loop = asyncio.get_running_loop() await loop.run_in_executor(None, connection.start) logger.debug(LogMessages.WS_CONNECTED, extra={"hub": name}) @handle_errors("disconnect") async def disconnect(self: "ProjectXRealtimeClientProtocol") -> None: """ Disconnect from ProjectX Gateway hubs. Gracefully closes both user and market hub connections. Safe to call even if not connected. Clears connection flags but preserves callbacks and subscriptions for potential reconnection. Example: >>> # Graceful shutdown >>> await client.disconnect() >>> print("Disconnected from ProjectX Gateway") >>> >>> # Can reconnect later >>> if await client.connect(): ... # Previous subscriptions must be re-established ... await client.subscribe_user_updates() Side Effects: - Sets self.user_connected = False - Sets self.market_connected = False - Stops SignalR connections Note: Does not clear callbacks or subscription lists, allowing for reconnection with the same configuration. """ with LogContext( logger, operation="disconnect", account_id=self.account_id, ): logger.debug(LogMessages.WS_DISCONNECT) async with self._connection_lock: loop = asyncio.get_running_loop() if self.user_connection: await loop.run_in_executor(None, self.user_connection.stop) self.user_connected = False if self.market_connection: await loop.run_in_executor(None, self.market_connection.stop) self.market_connected = False logger.debug(LogMessages.WS_DISCONNECTED) # Connection event handlers def _on_user_hub_open(self: "ProjectXRealtimeClientProtocol") -> None: """ Handle user hub connection open. Called by SignalR when user hub connection is established. Sets connection flag and logs success. Side Effects: - Sets self.user_connected = True - Logs connection success """ self.user_connected = True self.user_hub_ready.set() self.logger.info("✅ User hub connected") def _on_user_hub_close(self: "ProjectXRealtimeClientProtocol") -> None: """ Handle user hub connection close. Called by SignalR when user hub connection is lost. Clears connection flag and logs warning. Side Effects: - Sets self.user_connected = False - Logs disconnection warning Note: Automatic reconnection will attempt based on configuration. """ self.user_connected = False self.user_hub_ready.clear() self.logger.warning("❌ User hub disconnected") def _on_market_hub_open(self: "ProjectXRealtimeClientProtocol") -> None: """ Handle market hub connection open. Called by SignalR when market hub connection is established. Sets connection flag and logs success. Side Effects: - Sets self.market_connected = True - Logs connection success """ self.market_connected = True self.market_hub_ready.set() self.logger.info("✅ Market hub connected") def _on_market_hub_close(self: "ProjectXRealtimeClientProtocol") -> None: """ Handle market hub connection close. Called by SignalR when market hub connection is lost. Clears connection flag and logs warning. Side Effects: - Sets self.market_connected = False - Logs disconnection warning Note: Automatic reconnection will attempt based on configuration. """ self.market_connected = False self.market_hub_ready.clear() self.logger.warning("❌ Market hub disconnected") def _on_connection_error( self: "ProjectXRealtimeClientProtocol", hub: str, error: Any ) -> None: """ Handle connection errors. Processes errors from SignalR connections. Filters out normal completion messages that SignalR sends as part of its protocol. Args: hub (str): Hub name ("user" or "market") error: Error object or message from SignalR Side Effects: - Increments connection error counter for real errors - Logs errors (excludes CompletionMessage) Note: SignalR CompletionMessage is not an error - it's a normal protocol message. """ # Check if this is a SignalR CompletionMessage (not an error) error_type = type(error).__name__ if "CompletionMessage" in error_type: # This is a normal SignalR protocol message, not an error logger.debug(f"SignalR completion message from {hub} hub: {error}") return # Log actual errors logger.error(LogMessages.WS_ERROR, extra={"hub": hub, "error": str(error)}) self.stats["connection_errors"] += 1 @handle_errors("update JWT token", reraise=False, default_return=False) async def update_jwt_token( self: "ProjectXRealtimeClientProtocol", new_jwt_token: str, timeout: float = 30.0, ) -> bool: """ Update JWT token and reconnect with new credentials. **CRITICAL FIX**: Implements deadlock prevention through timeout-based reconnection and connection state recovery mechanisms. Handles JWT token refresh for expired or updated tokens. Disconnects current connections, updates URLs with new token, and re-establishes all subscriptions. **Deadlock Prevention Features (v3.3.1)**: - Connection lock timeout prevents indefinite waiting - Automatic rollback to original state on failure - Connection state recovery preserves subscriptions - Comprehensive error handling with cleanup Args: new_jwt_token (str): New JWT authentication token from AsyncProjectX timeout (float): Maximum time in seconds to wait for reconnection (default: 30.0) Prevents deadlocks by ensuring operation completes within timeout Returns: bool: True if reconnection successful with new token Process: 1. Acquire connection lock with timeout (DEADLOCK PREVENTION) 2. Store original state for rollback (STATE RECOVERY) 3. Disconnect existing connections 4. Update token and connection URLs 5. Reset connection state 6. Reconnect to both hubs with timeout 7. Re-subscribe to user updates 8. Re-subscribe to previous market data 9. Implement connection state recovery on failure (ROLLBACK) **Safety Mechanisms**: - **Timeout Protection**: 30-second default prevents indefinite blocking - **State Recovery**: Original connection state restored on failure - **Subscription Preservation**: Market data subscriptions restored automatically - **Error Isolation**: Failures don't leave client in inconsistent state Example: >>> # Token refresh with deadlock prevention >>> async def refresh_connection(): ... # Get new token ... await project_x.authenticate() ... new_token = project_x.session_token ... # Update with timeout for deadlock prevention ... if await realtime_client.update_jwt_token(new_token, timeout=45.0): ... print("Reconnected with new token") ... else: ... print("Reconnection failed, original state recovered") >>> >>> # Production usage with error handling >>> try: ... success = await realtime_client.update_jwt_token(new_token) ... if success: ... logger.info("Token refresh successful") ... else: ... logger.error("Token refresh failed - check logs") ... except TimeoutError: ... logger.error("Token refresh timed out - deadlock prevented") Side Effects: - Disconnects and reconnects both hubs - Re-subscribes to all previous subscriptions - Updates internal token and URLs - Implements recovery mechanism on failure **Performance Impact**: - Brief data gap during reconnection (~2-5 seconds) - Timeout overhead minimal for successful operations - State recovery adds safety with minimal performance cost Note: - Callbacks are preserved during reconnection - Market data subscriptions are restored automatically - **NEW**: Deadlock prevention eliminates indefinite blocking - **NEW**: Connection state recovery prevents inconsistent states """ with LogContext( logger, operation="update_jwt_token", account_id=self.account_id, timeout=timeout, ): logger.debug(LogMessages.AUTH_REFRESH) # Store original state for recovery original_token = self.jwt_token original_setup_complete = self.setup_complete original_subscriptions = list(self._subscribed_contracts) try: # Acquire connection lock with timeout to prevent deadlock async with asyncio.timeout(timeout): async with self._connection_lock: # Disconnect existing connections await self.disconnect() # Update JWT token self.jwt_token = new_jwt_token # Reset setup flag to force new connection setup self.setup_complete = False # Reconnect with timeout reconnect_success = False try: async with asyncio.timeout( timeout * 0.7 ): # Reserve time for subscriptions reconnect_success = await self.connect() except TimeoutError: logger.error( LogMessages.WS_ERROR, extra={ "error": f"Connection timeout after {timeout * 0.7}s" }, ) reconnect_success = False if reconnect_success: # Re-subscribe to user updates with timeout try: async with asyncio.timeout( timeout * 0.15 ): # Small portion for user updates await self.subscribe_user_updates() except TimeoutError: logger.warning( "User subscription timeout during token refresh" ) # Re-subscribe to market data with timeout if original_subscriptions: try: async with asyncio.timeout( timeout * 0.15 ): # Small portion for market data await self.subscribe_market_data( original_subscriptions ) except TimeoutError: logger.warning( "Market subscription timeout during token refresh" ) logger.debug(LogMessages.WS_RECONNECT) return True else: # Connection failed - initiate recovery logger.error( LogMessages.WS_ERROR, extra={ "error": "Failed to reconnect with new JWT token" }, ) await self._recover_connection_state( original_token, original_setup_complete, original_subscriptions, ) return False except TimeoutError: logger.error( LogMessages.WS_ERROR, extra={"error": f"Token refresh timeout after {timeout}s"}, ) # Attempt recovery on timeout await self._recover_connection_state( original_token, original_setup_complete, original_subscriptions ) return False except Exception as e: logger.error( LogMessages.WS_ERROR, extra={"error": f"Token refresh failed: {e}"}, ) # Attempt recovery on any other error await self._recover_connection_state( original_token, original_setup_complete, original_subscriptions ) return False async def _recover_connection_state( self: "ProjectXRealtimeClientProtocol", original_token: str, original_setup_complete: bool, original_subscriptions: list[str], ) -> None: """ Recover connection state after failed token refresh. Attempts to restore the original connection state when token refresh fails. This prevents the client from being left in an inconsistent state. Args: original_token: Original JWT token to restore original_setup_complete: Original setup completion state original_subscriptions: List of original market data subscriptions """ logger.info("Attempting connection state recovery after failed token refresh") try: # Restore original token self.jwt_token = original_token self.setup_complete = original_setup_complete # Clear any partial connection state self.user_connected = False self.market_connected = False self.user_hub_ready.clear() self.market_hub_ready.clear() # Try to reconnect with original token (short timeout) recovery_timeout = 10.0 try: async with asyncio.timeout(recovery_timeout): if await self.connect(): logger.info( "Successfully recovered connection with original token" ) # Restore subscriptions try: await self.subscribe_user_updates() if original_subscriptions: await self.subscribe_market_data(original_subscriptions) logger.info("Successfully restored subscriptions") except Exception as e: logger.warning(f"Failed to restore subscriptions: {e}") else: logger.error("Failed to recover connection state") # Mark as disconnected state self.user_connected = False self.market_connected = False except TimeoutError: logger.error(f"Connection recovery timeout after {recovery_timeout}s") # Mark as disconnected state self.user_connected = False self.market_connected = False except Exception as e: logger.error(f"Error during connection state recovery: {e}") # Ensure we're in a clean disconnected state self.user_connected = False self.market_connected = False self.user_hub_ready.clear() self.market_hub_ready.clear() def is_connected(self: "ProjectXRealtimeClientProtocol") -> bool: """ Check if both hubs are connected. Returns: bool: True only if both user and market hubs are connected Example: >>> if client.is_connected(): ... print("Fully connected") ... elif client.user_connected: ... print("Only user hub connected") ... elif client.market_connected: ... print("Only market hub connected") ... else: ... print("Not connected") Note: Both hubs must be connected for full functionality. Check individual flags for partial connection status. """ return self.user_connected and self.market_connected def get_stats(self: "ProjectXRealtimeClientProtocol") -> dict[str, Any]: """ Get connection statistics. Provides comprehensive statistics about connection health, event flow, and subscription status. Returns: dict[str, Any]: Statistics dictionary containing: - events_received (int): Total events processed - connection_errors (int): Total connection errors - last_event_time (datetime): Most recent event timestamp - connected_time (datetime): When connection established - user_connected (bool): User hub connection status - market_connected (bool): Market hub connection status - subscribed_contracts (int): Number of market subscriptions Example: >>> stats = client.get_stats() >>> print(f"Events received: {stats['events_received']}") >>> print(f"Uptime: {datetime.now() - stats['connected_time']}") >>> if stats["connection_errors"] > 10: ... print("Warning: High error count") >>> # Monitor event flow >>> last_event = stats["last_event_time"] >>> if last_event and (datetime.now() - last_event).seconds > 60: ... print("Warning: No events for 60 seconds") Use Cases: - Connection health monitoring - Debugging event flow issues - Uptime tracking - Error rate monitoring """ # Get task statistics from task manager (if available) task_stats = {} if hasattr(self, "get_task_stats"): task_stats = self.get_task_stats() return { **self.stats, "user_connected": self.user_connected, "market_connected": self.market_connected, "subscribed_contracts": len(self._subscribed_contracts), "task_stats": task_stats, }