griptape-nodes 0.70.1__py3-none-any.whl → 0.71.0__py3-none-any.whl

griptape_nodes/bootstrap/workflow_executors/local_session_workflow_executor.py

@@ -1,13 +1,10 @@
 from __future__ import annotations
 
 import asyncio
-import json
 import logging
-import threading
 from typing import TYPE_CHECKING, Any, Self
 
-from griptape_nodes.api_client import Client
-from griptape_nodes.app.app import WebSocketMessage
+from griptape_nodes.bootstrap.utils.subprocess_websocket_sender import SubprocessWebSocketSenderMixin
 from griptape_nodes.bootstrap.workflow_executors.local_workflow_executor import (
     LocalExecutorError,
     LocalWorkflowExecutor,
@@ -20,10 +17,12 @@ from griptape_nodes.retained_mode.events.base_events import (
     EventResultSuccess,
     ExecutionEvent,
     ExecutionGriptapeNodeEvent,
+    ProgressEvent,
     ResultPayload,
 )
 from griptape_nodes.retained_mode.events.execution_events import (
     ControlFlowCancelledEvent,
+    GriptapeEvent,
     StartFlowRequest,
     StartFlowResultFailure,
 )
@@ -36,7 +35,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
+class LocalSessionWorkflowExecutor(LocalWorkflowExecutor, SubprocessWebSocketSenderMixin):
     def __init__(
         self,
         session_id: str,
@@ -44,13 +43,8 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
         on_start_flow_result: Callable[[ResultPayload], None] | None = None,
     ):
         super().__init__(storage_backend=storage_backend)
-        self._session_id = session_id
+        self._init_websocket_sender(session_id)
         self._on_start_flow_result = on_start_flow_result
-        self._websocket_thread: threading.Thread | None = None
-        self._websocket_event_loop: asyncio.AbstractEventLoop | None = None
-        self._websocket_event_loop_ready = threading.Event()
-        self._ws_outgoing_queue: asyncio.Queue | None = None
-        self._shutdown_event: asyncio.Event | None = None
 
     async def __aenter__(self) -> Self:
         """Async context manager entry: initialize queue and broadcast app initialization."""
@@ -71,37 +65,12 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
         exc_tb: TracebackType | None,
     ) -> None:
         """Async context manager exit."""
-        self._stop_websocket_thread()
+        await self._stop_websocket_connection()
 
         GriptapeNodes.SessionManager().remove_session(self._session_id)
 
         # TODO: Broadcast shutdown https://github.com/griptape-ai/griptape-nodes/issues/2149
 
-    def _stop_websocket_thread(self) -> None:
-        """Stop the websocket thread."""
-        if self._websocket_thread is None or not self._websocket_thread.is_alive():
-            logger.debug("No websocket thread to stop")
-            return
-
-        logger.debug("Stopping websocket thread")
-        self._websocket_event_loop_ready.clear()
-
-        # Signal shutdown to the websocket tasks
-        if self._websocket_event_loop and self._websocket_event_loop.is_running() and self._shutdown_event:
-
-            def signal_shutdown() -> None:
-                if self._shutdown_event:
-                    self._shutdown_event.set()
-
-            self._websocket_event_loop.call_soon_threadsafe(signal_shutdown)
-
-        # Wait for thread to finish
-        self._websocket_thread.join(timeout=5.0)
-        if self._websocket_thread.is_alive():
-            logger.warning("Websocket thread did not stop gracefully")
-        else:
-            logger.info("Websocket thread stopped successfully")
-
     async def _process_execution_event_async(self, event: ExecutionGriptapeNodeEvent) -> None:
         """Process execution events asynchronously for real-time websocket emission."""
         logger.debug("REAL-TIME: Processing execution event for session %s", self._session_id)
@@ -144,7 +113,7 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
             await asyncio.sleep(1)
             raise LocalExecutorError(msg) from e
         finally:
-            self._stop_websocket_thread()
+            await self._stop_websocket_connection()
 
     async def _arun(  # noqa: C901, PLR0915
         self,
@@ -234,6 +203,16 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
                     background_tasks.add(task)
                     task.add_done_callback(_handle_task_done)
                     is_flow_finished, error = await self._handle_execution_event(event, flow_name)
+                elif isinstance(event, ProgressEvent):
+                    # Convert ProgressEvent to GriptapeEvent and emit via WebSocket
+                    payload = GriptapeEvent(
+                        node_name=event.node_name,
+                        parameter_name=event.parameter_name,
+                        type=type(event).__name__,
+                        value=event.value,
+                    )
+                    execution_event = ExecutionEvent(payload=payload)
+                    self.send_event("execution_event", execution_event.json())
 
                 event_queue.task_done()
 
@@ -251,135 +230,3 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
 
         if error is not None:
             raise error
-
-    async def _start_websocket_connection(self) -> None:
-        """Start websocket connection in a background thread for event emission."""
-        logger.info("Starting websocket connection for session %s", self._session_id)
-        self._websocket_thread = threading.Thread(target=self._start_websocket_thread, daemon=True)
-        self._websocket_thread.start()
-
-        if self._websocket_event_loop_ready.wait(timeout=10):
-            logger.info("Websocket thread ready")
-            await asyncio.sleep(1)  # Brief wait for connection to establish
-        else:
-            logger.error("Timeout waiting for websocket thread to start")
-
-    def _start_websocket_thread(self) -> None:
-        """Run WebSocket tasks in a separate thread with its own async loop."""
-        try:
-            # Create a new event loop for this thread
-            loop = asyncio.new_event_loop()
-            self._websocket_event_loop = loop
-            asyncio.set_event_loop(loop)
-
-            # Create the outgoing queue and shutdown event
-            self._ws_outgoing_queue = asyncio.Queue()
-            self._shutdown_event = asyncio.Event()
-
-            # Signal that websocket_event_loop is ready
-            self._websocket_event_loop_ready.set()
-            logger.info("Websocket thread started and ready")
-
-            # Run the async WebSocket tasks
-            loop.run_until_complete(self._run_websocket_tasks())
-        except Exception as e:
-            logger.error("WebSocket thread error: %s", e)
-        finally:
-            self._websocket_event_loop = None
-            self._websocket_event_loop_ready.clear()
-            self._shutdown_event = None
-            logger.info("Websocket thread ended")
-
-    async def _run_websocket_tasks(self) -> None:
-        """Run websocket tasks - establish connection and handle outgoing messages."""
-        logger.info("Creating Client for session %s", self._session_id)
-
-        async with Client() as client:
-            logger.info("WebSocket connection established for session %s", self._session_id)
-
-            try:
-                await self._send_outgoing_messages(client)
-            except Exception:
-                logger.exception("WebSocket tasks failed")
-            finally:
-                logger.info("WebSocket connection loop ended")
-
-    async def _send_outgoing_messages(self, client: Client) -> None:
-        """Send outgoing WebSocket messages from queue - matches app.py pattern exactly."""
-        if self._ws_outgoing_queue is None:
-            logger.error("No outgoing queue available")
-            return
-
-        logger.debug("Starting outgoing WebSocket request sender")
-
-        while True:
-            # Check if shutdown was requested
-            if self._shutdown_event and self._shutdown_event.is_set():
-                logger.info("Shutdown requested, ending message sender")
-                break
-
-            try:
-                # Get message from outgoing queue with timeout to allow shutdown checks
-                message = await asyncio.wait_for(self._ws_outgoing_queue.get(), timeout=1.0)
-            except TimeoutError:
-                # No message in queue, continue to check for shutdown
-                continue
-
-            try:
-                if isinstance(message, WebSocketMessage):
-                    topic = message.topic if message.topic else f"sessions/{self._session_id}/response"
-                    payload_dict = json.loads(message.payload)
-                    await client.publish(message.event_type, payload_dict, topic)
-                    logger.debug("DELIVERED: %s event", message.event_type)
-                else:
-                    logger.warning("Unknown outgoing message type: %s", type(message))
-            except Exception as e:
-                logger.error("Error sending outgoing WebSocket request: %s", e)
-            finally:
-                self._ws_outgoing_queue.task_done()
-
-    def send_event(self, event_type: str, payload: str) -> None:
-        """Send an event via websocket if connected - thread-safe version."""
-        # Wait for websocket event loop to be ready
-        if not self._websocket_event_loop_ready.wait(timeout=1.0):
-            logger.debug("Websocket not ready, event not sent: %s", event_type)
-            return
-
-        # Use run_coroutine_threadsafe to put message into WebSocket background thread queue
-        if self._websocket_event_loop is None:
-            logger.debug("WebSocket event loop not available for message: %s", event_type)
-            return
-
-        topic = f"sessions/{self._session_id}/response"
-        message = WebSocketMessage(event_type, payload, topic)
-
-        if self._ws_outgoing_queue is None:
-            logger.debug("No websocket queue available for event: %s", event_type)
-            return
-
-        try:
-            asyncio.run_coroutine_threadsafe(self._ws_outgoing_queue.put(message), self._websocket_event_loop)
-            logger.debug("SENT: %s event via websocket", event_type)
-        except Exception as e:
-            logger.error("Failed to queue event %s: %s", event_type, e)
-
-    async def _wait_for_websocket_queue_flush(self, timeout_seconds: float = 5.0) -> None:
-        """Wait for all websocket messages to be sent."""
-        if self._ws_outgoing_queue is None or self._websocket_event_loop is None:
-            return
-
-        async def _check_queue_empty() -> bool:
-            return self._ws_outgoing_queue.empty() if self._ws_outgoing_queue else True
-
-        start_time = asyncio.get_event_loop().time()
-        while asyncio.get_event_loop().time() - start_time < timeout_seconds:
-            future = asyncio.run_coroutine_threadsafe(_check_queue_empty(), self._websocket_event_loop)
-            try:
-                is_empty = future.result(timeout=0.1)
-                if is_empty:
-                    return
-            except Exception as e:
-                logger.debug("Error checking queue status: %s", e)
-            await asyncio.sleep(0.1)
-
-        logger.warning("Timeout waiting for websocket queue to flush")

griptape_nodes/bootstrap/workflow_executors/subprocess_workflow_executor.py

@@ -1,19 +1,16 @@
 from __future__ import annotations
 
-import asyncio
 import json
 import logging
 import tempfile
-import threading
-import uuid
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Self
 
 import anyio
 
-from griptape_nodes.api_client import Client
 from griptape_nodes.bootstrap.utils.python_subprocess_executor import PythonSubprocessExecutor
-from griptape_nodes.bootstrap.workflow_executors.local_session_workflow_executor import LocalSessionWorkflowExecutor
+from griptape_nodes.bootstrap.utils.subprocess_websocket_listener import SubprocessWebSocketListenerMixin
+from griptape_nodes.bootstrap.workflow_executors.workflow_executor import WorkflowExecutor
 from griptape_nodes.drivers.storage import StorageBackend
 from griptape_nodes.retained_mode.events.base_events import (
     EventResultFailure,
@@ -39,28 +36,23 @@ class SubprocessWorkflowExecutorError(Exception):
     """Exception raised during subprocess workflow execution."""
 
 
-class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessExecutor):
+class SubprocessWorkflowExecutor(WorkflowExecutor, PythonSubprocessExecutor, SubprocessWebSocketListenerMixin):
     def __init__(
         self,
         workflow_path: str,
         on_start_flow_result: Callable[[ResultPayload], None] | None = None,
+        on_event: Callable[[dict], None] | None = None,
         session_id: str | None = None,
     ) -> None:
+        WorkflowExecutor.__init__(self)
         PythonSubprocessExecutor.__init__(self)
+        self._init_websocket_listener(session_id=session_id, on_event=on_event)
         self._workflow_path = workflow_path
         self._on_start_flow_result = on_start_flow_result
-        # Generate a unique session ID for this execution
-        self._session_id = session_id or uuid.uuid4().hex
-        self._websocket_thread: threading.Thread | None = None
-        self._websocket_event_loop: asyncio.AbstractEventLoop | None = None
-        self._websocket_event_loop_ready = threading.Event()
-        self._event_handlers: dict[str, list] = {}
-        self._shutdown_event: asyncio.Event | None = None
         self._stored_exception: SubprocessWorkflowExecutorError | None = None
 
     async def __aenter__(self) -> Self:
-        """Async context manager entry: start WebSocket connection."""
-        logger.info("Starting WebSocket listener for session %s", self._session_id)
+        """Async context manager entry: start WebSocket listener."""
         await self._start_websocket_listener()
         return self
 
@@ -70,9 +62,8 @@ class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessE
         exc_val: BaseException | None,
         exc_tb: TracebackType | None,
     ) -> None:
-        """Async context manager exit: stop WebSocket connection."""
-        logger.info("Stopping WebSocket listener for session %s", self._session_id)
-        self._stop_websocket_listener()
+        """Async context manager exit: stop WebSocket listener."""
+        await self._stop_websocket_listener()
 
     async def arun(
         self,
@@ -138,106 +129,11 @@ class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessE
                 if self._stored_exception:
                     raise self._stored_exception
 
-    async def _start_websocket_listener(self) -> None:
-        """Start WebSocket connection to listen for events from the subprocess."""
-        logger.info("Starting WebSocket listener for session %s", self._session_id)
-        self._websocket_thread = threading.Thread(target=self._start_websocket_thread, daemon=True)
-        self._websocket_thread.start()
+    async def _handle_subprocess_event(self, event: dict) -> None:
+        """Handle executor-specific events from the subprocess.
 
-        if self._websocket_event_loop_ready.wait(timeout=10):
-            logger.info("WebSocket listener thread ready")
-        else:
-            logger.error("Timeout waiting for WebSocket listener thread to start")
-
-    def _stop_websocket_listener(self) -> None:
-        """Stop the WebSocket listener thread."""
-        if self._websocket_thread is None or not self._websocket_thread.is_alive():
-            return
-
-        logger.info("Stopping WebSocket listener thread")
-        self._websocket_event_loop_ready.clear()
-
-        # Signal shutdown to the websocket tasks
-        if self._websocket_event_loop and self._websocket_event_loop.is_running() and self._shutdown_event:
-
-            def signal_shutdown() -> None:
-                if self._shutdown_event:
-                    self._shutdown_event.set()
-
-            self._websocket_event_loop.call_soon_threadsafe(signal_shutdown)
-
-        # Wait for thread to finish
-        self._websocket_thread.join(timeout=5.0)
-        if self._websocket_thread.is_alive():
-            logger.warning("WebSocket listener thread did not stop gracefully")
-        else:
-            logger.info("WebSocket listener thread stopped successfully")
-
-    def _start_websocket_thread(self) -> None:
-        """Run WebSocket tasks in a separate thread with its own async loop."""
-        try:
-            # Create a new event loop for this thread
-            loop = asyncio.new_event_loop()
-            self._websocket_event_loop = loop
-            asyncio.set_event_loop(loop)
-
-            # Create shutdown event
-            self._shutdown_event = asyncio.Event()
-
-            # Signal that websocket_event_loop is ready
-            self._websocket_event_loop_ready.set()
-            logger.info("WebSocket listener thread started and ready")
-
-            # Run the async WebSocket listener
-            loop.run_until_complete(self._run_websocket_listener())
-        except Exception as e:
-            logger.error("WebSocket listener thread error: %s", e)
-        finally:
-            self._websocket_event_loop = None
-            self._websocket_event_loop_ready.clear()
-            self._shutdown_event = None
-            logger.info("WebSocket listener thread ended")
-
-    async def _run_websocket_listener(self) -> None:
-        """Run WebSocket listener - establish connection and handle incoming messages."""
-        logger.info("Creating Client for listening on session %s", self._session_id)
-
-        async with Client() as client:
-            logger.info("WebSocket connection established for session %s", self._session_id)
-
-            try:
-                await self._listen_for_messages(client)
-            except Exception:
-                logger.exception("WebSocket listener failed")
-            finally:
-                logger.info("WebSocket listener connection loop ended")
-
-    async def _listen_for_messages(self, client: Client) -> None:
-        """Listen for incoming WebSocket messages from the subprocess."""
-        logger.info("Starting to listen for WebSocket messages")
-
-        topic = f"sessions/{self._session_id}/response"
-        await client.subscribe(topic)
-
-        try:
-            async for message in client.messages:
-                if self._shutdown_event and self._shutdown_event.is_set():
-                    logger.info("Shutdown requested, ending message listener")
-                    break
-
-                try:
-                    logger.debug("Received WebSocket message: %s", message.get("type"))
-                    await self._process_event(message)
-
-                except Exception:
-                    logger.exception("Error processing WebSocket message")
-
-        except Exception as e:
-            logger.error("Error in WebSocket message listener: %s", e)
-            raise
-
-    async def _process_event(self, event: dict) -> None:
-        """Process events received from the subprocess via WebSocket."""
+        Processes execution events and result events.
+        """
         event_type = event.get("type", "unknown")
         if event_type == "execution_event":
             await self._process_execution_event(event)

griptape_nodes/bootstrap/workflow_publishers/local_session_workflow_publisher.py

@@ -0,0 +1,206 @@
+"""Local session workflow publisher with WebSocket event emission support.
+
+This module provides a workflow publisher that emits events over WebSocket
+for real-time progress updates during the publishing process.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Any, Self
+
+from griptape_nodes.bootstrap.utils.subprocess_websocket_sender import SubprocessWebSocketSenderMixin
+from griptape_nodes.bootstrap.workflow_publishers.local_workflow_publisher import (
+    LocalPublisherError,
+    LocalWorkflowPublisher,
+)
+from griptape_nodes.retained_mode.events.app_events import AppInitializationComplete
+from griptape_nodes.retained_mode.events.base_events import (
+    EventResultFailure,
+    EventResultSuccess,
+    ExecutionEvent,
+    ExecutionGriptapeNodeEvent,
+    ResultPayload,
+)
+from griptape_nodes.retained_mode.events.workflow_events import (
+    PublishWorkflowProgressEvent,
+    PublishWorkflowRequest,
+    PublishWorkflowResultFailure,
+)
+from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+logger = logging.getLogger(__name__)
+
+
+class LocalSessionWorkflowPublisher(LocalWorkflowPublisher, SubprocessWebSocketSenderMixin):
+    """Publisher with WebSocket support for sending events to parent process.
+
+    This publisher is used inside subprocesses to emit publishing progress events
+    over WebSocket back to the parent process.
+    """
+
+    def __init__(self, session_id: str) -> None:
+        super().__init__()
+        self._init_websocket_sender(session_id)
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry: initialize queue and start WebSocket connection."""
+        GriptapeNodes.EventManager().initialize_queue()
+        await GriptapeNodes.EventManager().broadcast_app_event(AppInitializationComplete())
+
+        logger.info("Setting up publishing session %s", self._session_id)
+        GriptapeNodes.SessionManager().save_session(self._session_id)
+        GriptapeNodes.SessionManager().active_session_id = self._session_id
+        await self._start_websocket_connection()
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Async context manager exit."""
+        await self._stop_websocket_connection()
+        GriptapeNodes.SessionManager().remove_session(self._session_id)
+
+    async def arun(
+        self,
+        workflow_name: str,
+        workflow_path: str,
+        publisher_name: str,
+        published_workflow_file_name: str,
+        **kwargs: Any,
+    ) -> None:
+        """Run the publish operation with WebSocket event emission enabled.
+
+        Progress events are emitted by the Library handling the publishing process,
+        which uses the GriptapeNodes event infrastructure to emit events that will
+        be sent over WebSocket to the parent process.
+        """
+        try:
+            await self._arun(
+                workflow_name=workflow_name,
+                workflow_path=workflow_path,
+                publisher_name=publisher_name,
+                published_workflow_file_name=published_workflow_file_name,
+                **kwargs,
+            )
+        except Exception as e:
+            msg = f"Unexpected error during publish: {e}"
+            logger.exception(msg)
+            raise LocalPublisherError(msg) from e
+        finally:
+            await self._stop_websocket_connection()
+
+    async def _arun(  # noqa: C901, PLR0915
+        self,
+        workflow_name: str,
+        workflow_path: str,
+        publisher_name: str,
+        published_workflow_file_name: str,
+        **kwargs: Any,
+    ) -> None:
+        """Internal async run method with event queue monitoring and websocket integration."""
+        # Load the workflow into memory
+        await self.aprepare_workflow_for_run(flow_input={}, workflow_path=workflow_path)
+
+        pickle_control_flow_result = kwargs.get("pickle_control_flow_result", False)
+        publish_workflow_request = PublishWorkflowRequest(
+            workflow_name=workflow_name,
+            publisher_name=publisher_name,
+            published_workflow_file_name=published_workflow_file_name,
+            pickle_control_flow_result=pickle_control_flow_result,
+        )
+
+        # Send the publish request async (fire and forget pattern)
+        publish_task = asyncio.create_task(GriptapeNodes.ahandle_request(publish_workflow_request))
+
+        is_publish_finished = False
+        error: Exception | None = None
+        background_tasks: set[asyncio.Task] = set()
+
+        def _handle_publish_result(task: asyncio.Task[ResultPayload]) -> None:
+            nonlocal is_publish_finished, error
+            try:
+                publish_result = task.result()
+
+                if isinstance(publish_result, PublishWorkflowResultFailure):
+                    msg = f"Failed to publish workflow: {publish_result.result_details}"
+                    logger.error(msg)
+                    event_result_failure = EventResultFailure(request=publish_workflow_request, result=publish_result)
+                    self.send_event("failure_result", event_result_failure.json())
+                    is_publish_finished = True
+                    error = LocalPublisherError(msg)
+                else:
+                    logger.info("Published workflow successfully")
+                    event_result_success = EventResultSuccess(request=publish_workflow_request, result=publish_result)
+                    self.send_event("success_result", event_result_success.json())
+                    is_publish_finished = True
+
+            except Exception as e:
+                msg = "Error during publish workflow"
+                logger.exception(msg)
+                is_publish_finished = True
+                error = e
+                # Add a dummy event to wake up the loop in failure cases
+                event_queue = GriptapeNodes.EventManager().event_queue
+                queue_event_task = asyncio.create_task(event_queue.put(None))
+                background_tasks.add(queue_event_task)
+                queue_event_task.add_done_callback(background_tasks.discard)
+
+        publish_task.add_done_callback(_handle_publish_result)
+
+        logger.info("Publish workflow request sent! Processing events...")
+
+        def _handle_task_done(task: asyncio.Task) -> None:
+            background_tasks.discard(task)
+            if task.exception() and not task.cancelled():
+                logger.exception("Background task failed", exc_info=task.exception())
+
+        event_queue = GriptapeNodes.EventManager().event_queue
+        while not is_publish_finished:
+            try:
+                event = await event_queue.get()
+
+                # Handle the dummy wake up event (None)
+                if event is None:
+                    event_queue.task_done()
+                    continue
+
+                logger.debug("Processing publish event: %s", type(event).__name__)
+
+                if isinstance(event, ExecutionGriptapeNodeEvent):
+                    # Unwrap the event and check if it contains a PublishWorkflowProgressEvent
+                    wrapped_event = event.wrapped_event
+                    if isinstance(wrapped_event, ExecutionEvent) and isinstance(
+                        wrapped_event.payload, PublishWorkflowProgressEvent
+                    ):
+                        self.send_event("execution_event", wrapped_event.json())
+                        logger.debug(
+                            "Emitted progress event: %.1f%% - %s",
+                            wrapped_event.payload.progress,
+                            wrapped_event.payload.message,
+                        )
+
+                event_queue.task_done()
+
+            except Exception as e:
+                msg = f"Error handling queue event: {e}"
+                logger.exception(msg)
+                error = LocalPublisherError(msg)
+                break
+
+        if background_tasks:
+            logger.info("Waiting for %d background tasks to complete", len(background_tasks))
+            await asyncio.gather(*background_tasks, return_exceptions=True)
+
+        await self._wait_for_websocket_queue_flush()
+
+        if error is not None:
+            raise error