rest adaper fix

Author: marcomelloni

simulation_bridge/src/protocol_adapters/rest/rest_adapter.py

@@ -1,5 +1,5 @@
 from quart import Quart, request, Response
-from hypercorn.config import Config
+from hypercorn.config import Config as HyperConfig
 from hypercorn.asyncio import serve
 import asyncio
 import yaml
@@ -17,22 +17,36 @@ class RESTAdapter(ProtocolAdapter):
     """REST protocol adapter implementation using Quart and Hypercorn."""
 
     def _get_config(self) -> Dict[str, Any]:
+        """Get REST configuration from config manager."""
         return self.config_manager.get_rest_config()
 
     def __init__(self, config_manager: ConfigManager):
+        """Initialize REST adapter with configuration.
+
+        Args:
+            config_manager: Configuration manager instance
+        """
         super().__init__(config_manager)
         self.app = Quart(__name__)
         self._setup_routes()
-        self._active_streams = {}  # client_id -> asyncio.Queue
+        self.server = None
+        self._active_streams = {}  # Store active streams by client_id
+        # Main event loop
+        self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._running = False
-        self._server_task: Optional[asyncio.Task] = None
         logger.debug("REST - Adapter initialized with config: host=%s, port=%s",
                      self.config['host'], self.config['port'])
 
     def _setup_routes(self) -> None:
+        """Set up the streaming endpoint."""
         self.app.post(self.config['endpoint'])(self._handle_streaming_message)
 
     async def _handle_streaming_message(self) -> Response:
+        """Handle incoming messages with streaming response.
+
+        Returns:
+            Response: Streaming response with simulation results
+        """
         content_type = request.headers.get('content-type', '')
         body = await request.get_data()
 
@@ -58,6 +72,7 @@ async def _handle_streaming_message(self) -> Response:
         producer = simulation.get('client_id', 'unknown')
         consumer = simulation.get('simulator', 'unknown')
 
+        # Add bridge metadata
         message['bridge_meta'] = {
             'protocol': 'rest',
             'producer': producer,
@@ -67,14 +82,15 @@ async def _handle_streaming_message(self) -> Response:
         logger.debug(
             "REST - Processing message from producer: %s, simulator: %s",
             producer, consumer)
-
+        # Use SignalManager to send the signal
         signal('message_received_input_rest').send(
             message=message,
             producer=producer,
             consumer=consumer,
             protocol='rest'
         )
 
+        # Create a queue for this client's messages
         queue = asyncio.Queue()
         self._active_streams[producer] = queue
 
@@ -85,12 +101,23 @@ async def _handle_streaming_message(self) -> Response:
         )
 
     def _parse_message(self, body: bytes, content_type: str) -> Dict[str, Any]:
+        """Parse message body based on content type.
+
+        Args:
+            body: Raw message body
+            content_type: Content type header
+
+        Returns:
+            Dict[str, Any]: Parsed message
+        """
         if 'yaml' in content_type:
             logger.debug("REST - Attempting to parse message as YAML")
             return yaml.safe_load(body)
         elif 'json' in content_type:
             logger.debug("REST - Attempting to parse message as JSON")
             return json.loads(body)
+
+        # Fallback: try YAML, then JSON, then raw text
         try:
             logger.debug(
                 "REST - Attempting to parse message as YAML (fallback)")
@@ -108,95 +135,131 @@ def _parse_message(self, body: bytes, content_type: str) -> Dict[str, Any]:
                 }
 
     async def _generate_response(
-        self, producer: str, queue: asyncio.Queue
-    ) -> AsyncGenerator[str, None]:
+            self, producer: str, queue: asyncio.Queue) -> AsyncGenerator[str, None]:
+        """Generate streaming response.
+
+        Args:
+            producer: Client ID
+            queue: Message queue for this client
+
+        Yields:
+            str: JSON-encoded messages
+        """
         try:
+            # Send initial acknowledgment
             yield json.dumps({"status": "processing"}) + "\n"
+            # Keep the connection open and wait for results
             while True:
                 try:
                     result = await asyncio.wait_for(queue.get(), timeout=600)
                     yield json.dumps(result) + "\n"
                 except asyncio.TimeoutError:
-                    yield json.dumps({
-                        "status": "timeout",
-                        "error": "No response received within timeout"
-                    }) + "\n"
+                    yield json.dumps({"status": "timeout", "error": "No response received within timeout"}) + "\n"
                     break
                 except Exception as e:
                     logger.error("REST - Error in stream: %s", e)
                     yield json.dumps({"status": "error", "error": str(e)}) + "\n"
                     break
         finally:
-            self._active_streams.pop(producer, None)
+            # Clean up when the stream ends
+            if producer in self._active_streams:
+                del self._active_streams[producer]
 
     async def send_result(self, producer: str, result: Dict[str, Any]) -> None:
+        """Send a result message to a specific client.
+
+        Args:
+            producer: Client ID
+            result: Result message to send
+        """
         if producer in self._active_streams:
             await self._active_streams[producer].put(result)
         else:
             logger.warning(
-                "REST - No active stream found for producer: %s",
-                producer)
-
-    async def _start_server(self):
-        config = Config()
-        config.bind = [f"{self.config['host']}:{self.config['port']}"]
-        config.certfile = self.config.get('certfile')
-        config.keyfile = self.config.get('keyfile')
-        config.worker_class = "asyncio"
+                "REST - No active stream found for producer: %s", producer)
+
+    async def _start_server(self) -> None:
+        """Start the Hypercorn server."""
+        self._loop = asyncio.get_running_loop()  # Save main event loop
+
+        config = HyperConfig()
+        config.errorlog = logger  # Use the main logger for error logs
+        config.accesslog = logger  # Use the main logger for access logs
+        config.bind = ["%s:%s" % (self.config['host'], self.config['port'])]
         config.use_reloader = False
+        config.worker_class = "asyncio"
+        config.alpn_protocols = ["h2", "http/1.1"]
 
-        logger.debug("REST - Starting Hypercorn server")
+        if self.config.get('certfile') and self.config.get('keyfile'):
+            config.certfile = self.config['certfile']
+            config.keyfile = self.config['keyfile']
         await serve(self.app, config)
 
-    def start(self):
-        """Start the REST server in the current asyncio event loop."""
-        if self._running:
-            logger.warning("REST - Adapter already running")
-            return
-        self._running = True
+    def start(self) -> None:
+        """Start the REST server."""
+        logger.debug(
+            "REST - Starting adapter on %s:%s",
+            self.config['host'], self.config['port'])
         try:
-            loop = asyncio.get_event_loop()
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-
-        self._server_task = loop.create_task(self._start_server())
-        logger.debug("REST - Server started as asyncio task")
-
-    def stop(self):
-        """Stop the REST adapter."""
-        logger.debug("REST - Stopping adapter")
-        self._running = False
-        if self._server_task:
-            self._server_task.cancel()
+            asyncio.run(self._start_server())
+            self._running = True
+        except Exception as e:
+            logger.error("REST - Error starting server: %s", e)
+            raise
 
     def send_result_sync(self, producer: str, result: Dict[str, Any]) -> None:
+        """Synchronous wrapper for sending result messages.
+
+        Args:
+            producer: Client ID
+            result: Result message to send
+        """
         if producer not in self._active_streams:
             logger.warning(
-                "REST - No active stream found for producer: %s. Available streams: %s",
+                "REST - No active stream found for producer: %s. "
+                "Available streams: %s",
                 producer, list(self._active_streams.keys())
             )
             return
 
-        try:
-            loop = asyncio.get_event_loop()
-            if loop.is_running():
-                future = asyncio.run_coroutine_threadsafe(
-                    self.send_result(producer, result),
-                    loop
-                )
+        if self._loop and self._loop.is_running():
+            # Use run_coroutine_threadsafe to execute coroutine in main loop
+            future = asyncio.run_coroutine_threadsafe(
+                self.send_result(producer, result),
+                self._loop
+            )
+            try:
+                # Optional: wait for result with short timeout
                 future.result(timeout=5)
-            else:
-                logger.error(
-                    "REST - Event loop not running; cannot send result.")
-        except Exception as e:
-            logger.error("REST - Error sending result: %s", e)
+            except Exception as e:
+                logger.error("REST - Error sending result: %s", e)
+        else:
+            logger.error("REST - Event loop not running; cannot send result.")
+
+    def stop(self) -> None:
+        """Stop the REST server."""
+        logger.debug("REST - Stopping adapter")
+        self._running = False
+        if self.server:
+            self.server.close()
 
     def _handle_message(self, message: Dict[str, Any]) -> None:
-        # REST handled by Quart routes
+        """Handle incoming messages (required by ProtocolAdapter).
+
+        Args:
+            message: Message to handle
+        """
+        # For REST, this is handled by the Quart endpoint
         pass
 
-    def publish_result_message_rest(self, sender, **kwargs):
+    def publish_result_message_rest(self, sender, **kwargs):  # pylint: disable=unused-argument
+        """
+        Publish result message via REST adapter.
+
+        Args:
+            message: Message payload to send
+            destination: REST endpoint destination
+        """
         try:
             message = kwargs.get('message', {})
             destination = message.get('destinations', [])[0]