refactor: half way to process per session

core/communication/session_integration.py

@@ -0,0 +1,319 @@
+"""
+Integration layer between WebSocket handler and Session Process Manager.
+Bridges the existing WebSocket protocol with the new session-based architecture.
+"""
+
+import asyncio
+import logging
+from typing import Dict, Any, Optional
+import numpy as np
+
+from ..processes.session_manager import SessionProcessManager
+from ..processes.communication import DetectionResultResponse, ErrorResponse
+from .state import worker_state
+from .messages import serialize_outgoing_message
+# Streaming is now handled directly by session workers - no shared stream manager needed
+
+logger = logging.getLogger(__name__)
+
+
+class SessionWebSocketIntegration:
+    """
+    Integration layer that connects WebSocket protocol with Session Process Manager.
+    Maintains compatibility with existing WebSocket message handling.
+    """
+
+    def __init__(self, websocket_handler=None):
+        """
+        Initialize session WebSocket integration.
+
+        Args:
+            websocket_handler: Reference to WebSocket handler for sending messages
+        """
+        self.websocket_handler = websocket_handler
+        self.session_manager = SessionProcessManager()
+
+        # Track active subscriptions for compatibility
+        self.active_subscriptions: Dict[str, Dict[str, Any]] = {}
+
+        # Set up callbacks
+        self.session_manager.set_detection_result_callback(self._on_detection_result)
+        self.session_manager.set_error_callback(self._on_session_error)
+
+    async def start(self):
+        """Start the session integration."""
+        await self.session_manager.start()
+        logger.info("Session WebSocket integration started")
+
+    async def stop(self):
+        """Stop the session integration."""
+        await self.session_manager.stop()
+        logger.info("Session WebSocket integration stopped")
+
+    async def handle_set_subscription_list(self, message) -> bool:
+        """
+        Handle setSubscriptionList message by managing session processes.
+
+        Args:
+            message: SetSubscriptionListMessage
+
+        Returns:
+            True if successful
+        """
+        try:
+            logger.info(f"Processing subscription list with {len(message.subscriptions)} subscriptions")
+
+            new_subscription_ids = set()
+            for subscription in message.subscriptions:
+                subscription_id = subscription.subscriptionIdentifier
+                new_subscription_ids.add(subscription_id)
+
+                # Check if this is a new subscription
+                if subscription_id not in self.active_subscriptions:
+                    logger.info(f"Creating new session for subscription: {subscription_id}")
+
+                    # Convert subscription to configuration dict
+                    subscription_config = {
+                        'subscriptionIdentifier': subscription.subscriptionIdentifier,
+                        'rtspUrl': getattr(subscription, 'rtspUrl', None),
+                        'snapshotUrl': getattr(subscription, 'snapshotUrl', None),
+                        'snapshotInterval': getattr(subscription, 'snapshotInterval', 5000),
+                        'modelUrl': subscription.modelUrl,
+                        'modelId': subscription.modelId,
+                        'modelName': subscription.modelName,
+                        'cropX1': subscription.cropX1,
+                        'cropY1': subscription.cropY1,
+                        'cropX2': subscription.cropX2,
+                        'cropY2': subscription.cropY2
+                    }
+
+                    # Create session process
+                    success = await self.session_manager.create_session(
+                        subscription_id, subscription_config
+                    )
+
+                    if success:
+                        self.active_subscriptions[subscription_id] = subscription_config
+                        logger.info(f"Session created successfully for {subscription_id}")
+
+                        # Stream handling is now integrated into session worker process
+                    else:
+                        logger.error(f"Failed to create session for {subscription_id}")
+                        return False
+
+                else:
+                    # Update existing subscription configuration if needed
+                    self.active_subscriptions[subscription_id].update({
+                        'modelUrl': subscription.modelUrl,
+                        'modelId': subscription.modelId,
+                        'modelName': subscription.modelName,
+                        'cropX1': subscription.cropX1,
+                        'cropY1': subscription.cropY1,
+                        'cropX2': subscription.cropX2,
+                        'cropY2': subscription.cropY2
+                    })
+
+            # Remove sessions for subscriptions that are no longer active
+            current_subscription_ids = set(self.active_subscriptions.keys())
+            removed_subscriptions = current_subscription_ids - new_subscription_ids
+
+            for subscription_id in removed_subscriptions:
+                logger.info(f"Removing session for subscription: {subscription_id}")
+                await self.session_manager.remove_session(subscription_id)
+                del self.active_subscriptions[subscription_id]
+
+            # Update worker state for compatibility
+            worker_state.set_subscriptions(message.subscriptions)
+
+            logger.info(f"Subscription list processed: {len(new_subscription_ids)} active sessions")
+            return True
+
+        except Exception as e:
+            logger.error(f"Error handling subscription list: {e}", exc_info=True)
+            return False
+
+    async def handle_set_session_id(self, message) -> bool:
+        """
+        Handle setSessionId message by forwarding to appropriate session process.
+
+        Args:
+            message: SetSessionIdMessage
+
+        Returns:
+            True if successful
+        """
+        try:
+            display_id = message.payload.displayIdentifier
+            session_id = message.payload.sessionId
+
+            logger.info(f"Setting session ID {session_id} for display {display_id}")
+
+            # Find subscription identifier for this display
+            subscription_id = None
+            for sub_id in self.active_subscriptions.keys():
+                # Extract display identifier from subscription identifier
+                if display_id in sub_id:
+                    subscription_id = sub_id
+                    break
+
+            if not subscription_id:
+                logger.error(f"No active subscription found for display {display_id}")
+                return False
+
+            # Forward to session process
+            success = await self.session_manager.set_session_id(
+                subscription_id, str(session_id), display_id
+            )
+
+            if success:
+                # Update worker state for compatibility
+                worker_state.set_session_id(display_id, session_id)
+                logger.info(f"Session ID {session_id} set successfully for {display_id}")
+            else:
+                logger.error(f"Failed to set session ID {session_id} for {display_id}")
+
+            return success
+
+        except Exception as e:
+            logger.error(f"Error setting session ID: {e}", exc_info=True)
+            return False
+
+    async def process_frame(self, subscription_id: str, frame: np.ndarray, display_id: str, timestamp: float = None) -> bool:
+        """
+        Process frame through appropriate session process.
+
+        Args:
+            subscription_id: Subscription identifier
+            frame: Frame to process
+            display_id: Display identifier
+            timestamp: Frame timestamp
+
+        Returns:
+            True if frame was processed successfully
+        """
+        try:
+            if timestamp is None:
+                timestamp = asyncio.get_event_loop().time()
+
+            # Forward frame to session process
+            success = await self.session_manager.process_frame(
+                subscription_id, frame, display_id, timestamp
+            )
+
+            if not success:
+                logger.warning(f"Failed to process frame for subscription {subscription_id}")
+
+            return success
+
+        except Exception as e:
+            logger.error(f"Error processing frame for {subscription_id}: {e}", exc_info=True)
+            return False
+
+    async def _on_detection_result(self, subscription_id: str, response: DetectionResultResponse):
+        """
+        Handle detection result from session process.
+
+        Args:
+            subscription_id: Subscription identifier
+            response: Detection result response
+        """
+        try:
+            logger.debug(f"Received detection result from {subscription_id}: phase={response.phase}")
+
+            # Send imageDetection message via WebSocket (if needed)
+            if self.websocket_handler and hasattr(self.websocket_handler, 'send_message'):
+                from .models import ImageDetectionMessage, DetectionData
+
+                # Convert response detections to the expected format
+                # The DetectionData expects modelId and modelName, and detection dict
+                detection_data = DetectionData(
+                    detection=response.detections,
+                    modelId=getattr(response, 'model_id', 0),  # Get from response if available
+                    modelName=getattr(response, 'model_name', 'unknown')  # Get from response if available
+                )
+
+                # Convert timestamp to string format if it exists
+                timestamp_str = None
+                if hasattr(response, 'timestamp') and response.timestamp:
+                    from datetime import datetime
+                    if isinstance(response.timestamp, (int, float)):
+                        # Convert Unix timestamp to ISO format string
+                        timestamp_str = datetime.fromtimestamp(response.timestamp).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+                    else:
+                        timestamp_str = str(response.timestamp)
+
+                detection_message = ImageDetectionMessage(
+                    subscriptionIdentifier=subscription_id,
+                    data=detection_data,
+                    timestamp=timestamp_str
+                )
+
+                serialized = serialize_outgoing_message(detection_message)
+                await self.websocket_handler.send_message(serialized)
+
+        except Exception as e:
+            logger.error(f"Error handling detection result from {subscription_id}: {e}", exc_info=True)
+
+    async def _on_session_error(self, subscription_id: str, error_response: ErrorResponse):
+        """
+        Handle error from session process.
+
+        Args:
+            subscription_id: Subscription identifier
+            error_response: Error response
+        """
+        logger.error(f"Session error from {subscription_id}: {error_response.error_type} - {error_response.error_message}")
+
+        # Send error message via WebSocket if needed
+        if self.websocket_handler and hasattr(self.websocket_handler, 'send_message'):
+            error_message = {
+                'type': 'sessionError',
+                'payload': {
+                    'subscriptionIdentifier': subscription_id,
+                    'errorType': error_response.error_type,
+                    'errorMessage': error_response.error_message,
+                    'timestamp': error_response.timestamp
+                }
+            }
+
+            try:
+                serialized = serialize_outgoing_message(error_message)
+                await self.websocket_handler.send_message(serialized)
+            except Exception as e:
+                logger.error(f"Failed to send error message: {e}")
+
+    def get_session_stats(self) -> Dict[str, Any]:
+        """
+        Get statistics about active sessions.
+
+        Returns:
+            Dictionary with session statistics
+        """
+        return {
+            'active_sessions': self.session_manager.get_session_count(),
+            'max_sessions': self.session_manager.max_concurrent_sessions,
+            'subscriptions': list(self.active_subscriptions.keys())
+        }
+
+    async def handle_progression_stage(self, message) -> bool:
+        """
+        Handle setProgressionStage message.
+
+        Args:
+            message: SetProgressionStageMessage
+
+        Returns:
+            True if successful
+        """
+        try:
+            # For now, just update worker state for compatibility
+            # In future phases, this could be forwarded to session processes
+            worker_state.set_progression_stage(
+                message.payload.displayIdentifier,
+                message.payload.progressionStage
+            )
+            return True
+        except Exception as e:
+            logger.error(f"Error handling progression stage: {e}", exc_info=True)
+            return False
+