griptape-nodes 0.57.1__py3-none-any.whl → 0.58.1__py3-none-any.whl

griptape_nodes/app/app.py

@@ -3,20 +3,13 @@ from __future__ import annotations
 import asyncio
 import json
 import logging
-import os
-import sys
 import threading
 from dataclasses import dataclass
-from typing import Any
-from urllib.parse import urljoin
 
-from rich.align import Align
 from rich.console import Console
 from rich.logging import RichHandler
-from rich.panel import Panel
-from websockets.asyncio.client import connect
-from websockets.exceptions import ConnectionClosed, ConnectionClosedError, WebSocketException
 
+from griptape_nodes.api_client import Client
 from griptape_nodes.retained_mode.events import app_events, execution_events
 
 # This import is necessary to register all events, even if not technically used
@@ -98,11 +91,15 @@ class EventLogHandler(logging.Handler):
 logger = logging.getLogger("griptape_nodes_app")
 
 griptape_nodes_logger = logging.getLogger("griptape_nodes")
+griptape_nodes_logger.addHandler(EventLogHandler())
+griptape_nodes_logger.setLevel(logging.INFO)
+
+# Root logger only gets RichHandler for console output
 logging.basicConfig(
     level=logging.INFO,
     format="%(message)s",
     datefmt="[%X]",
-    handlers=[EventLogHandler(), RichHandler(show_time=True, show_path=False, markup=True, rich_tracebacks=True)],
+    handlers=[RichHandler(show_time=True, show_path=False, markup=True, rich_tracebacks=True)],
 )
 
 console = Console()
@@ -120,16 +117,12 @@ def start_app() -> None:
 
 async def astart_app() -> None:
     """New async app entry point."""
-    api_key = _ensure_api_key()
-
     # Initialize event queue in main thread
     griptape_nodes.EventManager().initialize_queue()
 
     try:
         # Start WebSocket tasks in daemon thread
-        threading.Thread(
-            target=_start_websocket_connection, args=(api_key,), daemon=True, name="websocket-tasks"
-        ).start()
+        threading.Thread(target=_start_websocket_connection, daemon=True, name="websocket-tasks").start()
 
         # Run event processing on main thread
         await _process_event_queue()
@@ -139,7 +132,7 @@ async def astart_app() -> None:
         raise
 
 
-def _start_websocket_connection(api_key: str) -> None:
+def _start_websocket_connection() -> None:
     """Run WebSocket tasks in a separate thread with its own async loop."""
     global websocket_event_loop  # noqa: PLW0603
     try:
@@ -152,7 +145,7 @@ def _start_websocket_connection(api_key: str) -> None:
         websocket_event_loop_ready.set()
 
         # Run the async WebSocket tasks
-        loop.run_until_complete(_run_websocket_tasks(api_key))
+        loop.run_until_complete(_run_websocket_tasks())
     except Exception as e:
         logger.error("WebSocket thread error: %s", e)
         raise
@@ -161,80 +154,40 @@ def _start_websocket_connection(api_key: str) -> None:
         websocket_event_loop_ready.clear()
 
 
-async def _run_websocket_tasks(api_key: str) -> None:
+async def _run_websocket_tasks() -> None:
     """Run WebSocket tasks - async version."""
-    # Create WebSocket connection for this thread
-    connection_stream = _create_websocket_connection(api_key)
-
-    # Track if this is the first connection
-    initialized = False
-
-    async for ws_connection in connection_stream:
+    async with Client() as client:
         logger.debug("WebSocket connection established")
-        try:
-            # Emit initialization event only for the first connection
-            if not initialized:
-                griptape_nodes.EventManager().put_event(AppEvent(payload=app_events.AppInitializationComplete()))
-                initialized = True
-
-            # Emit connection established event for every connection
-            griptape_nodes.EventManager().put_event(AppEvent(payload=app_events.AppConnectionEstablished()))
-
-            async with asyncio.TaskGroup() as tg:
-                tg.create_task(_process_incoming_messages(ws_connection))
-                tg.create_task(_send_outgoing_messages(ws_connection))
-        except (ExceptionGroup, ConnectionClosed, ConnectionClosedError):
-            logger.info("WebSocket connection closed, reconnecting...")
-            continue
-        except Exception:
-            logger.exception("WebSocket tasks failed")
-            await asyncio.sleep(2.0)  # Wait before retry
-            continue
-
-
-def _ensure_api_key() -> str:
-    secrets_manager = griptape_nodes.SecretsManager()
-    api_key = secrets_manager.get_secret("GT_CLOUD_API_KEY")
-    if api_key is None:
-        message = Panel(
-            Align.center(
-                "[bold red]Nodes API key is not set, please run [code]gtn init[/code] with a valid key: [/bold red]"
-                "[code]gtn init --api-key <your key>[/code]\n"
-                "[bold red]You can generate a new key from [/bold red][bold blue][link=https://nodes.griptape.ai]https://nodes.griptape.ai[/link][/bold blue]",
-            ),
-            title="[red]X[/red] Missing Nodes API Key",
-            border_style="red",
-            padding=(1, 4),
-        )
-        console.print(message)
-        sys.exit(1)
 
-    return api_key
+        griptape_nodes.EventManager().put_event(AppEvent(payload=app_events.AppInitializationComplete()))
+        griptape_nodes.EventManager().put_event(AppEvent(payload=app_events.AppConnectionEstablished()))
 
+        async with asyncio.TaskGroup() as tg:
+            tg.create_task(_process_incoming_messages(client))
+            tg.create_task(_send_outgoing_messages(client))
 
-async def _process_incoming_messages(ws_connection: Any) -> None:
+
+async def _process_incoming_messages(client: Client) -> None:
     """Process incoming WebSocket requests from Nodes API."""
     logger.debug("Processing incoming WebSocket requests from WebSocket connection")
 
-    async for message in ws_connection:
-        try:
-            data = json.loads(message)
-            await _process_api_event(data)
-        except Exception:
-            logger.exception("Error processing event, skipping.")
+    topics = ["request"]
+    engine_id = griptape_nodes.get_engine_id()
+    if engine_id:
+        topics.append(f"engines/{engine_id}/request")
 
+    session_id = griptape_nodes.get_session_id()
+    if session_id:
+        topics.append(f"sessions/{session_id}/request")
 
-def _create_websocket_connection(api_key: str) -> Any:
-    """Create an async WebSocket connection to the Nodes API."""
-    endpoint = urljoin(
-        os.getenv("GRIPTAPE_NODES_API_BASE_URL", "https://api.nodes.griptape.ai").replace("http", "ws"),
-        "/ws/engines/events?version=v2",
-    )
+    for topic in topics:
+        await client.subscribe(topic)
 
-    return connect(
-        endpoint,
-        additional_headers={"Authorization": f"Bearer {api_key}"},
-    )
+    async for message in client.messages:
+        try:
+            await _process_api_event(message)
+        except Exception:
+            logger.exception("Error processing event, skipping.")
 
 
 async def _process_api_event(event: dict) -> None:
@@ -276,7 +229,7 @@ async def _process_api_event(event: dict) -> None:
         griptape_nodes.EventManager().put_event(request_event)
 
 
-async def _send_outgoing_messages(ws_connection: Any) -> None:
+async def _send_outgoing_messages(client: Client) -> None:
     """Send outgoing WebSocket requests from queue on background thread."""
     logger.debug("Starting outgoing WebSocket request sender")
 
@@ -286,11 +239,14 @@ async def _send_outgoing_messages(ws_connection: Any) -> None:
 
         try:
             if isinstance(message, WebSocketMessage):
-                await _send_websocket_message(ws_connection, message.event_type, message.payload, message.topic)
+                # Use client to publish message
+                topic = message.topic if message.topic else _determine_response_topic()
+                payload_dict = json.loads(message.payload)
+                await client.publish(message.event_type, payload_dict, topic)
             elif isinstance(message, SubscribeCommand):
-                await _send_subscribe_command(ws_connection, message.topic)
+                await client.subscribe(message.topic)
             elif isinstance(message, UnsubscribeCommand):
-                await _send_unsubscribe_command(ws_connection, message.topic)
+                await client.unsubscribe(message.topic)
             else:
                 logger.warning("Unknown outgoing message type: %s", type(message))
         except Exception as e:
@@ -299,44 +255,6 @@ async def _send_outgoing_messages(ws_connection: Any) -> None:
             ws_outgoing_queue.task_done()
 
 
-async def _send_websocket_message(ws_connection: Any, event_type: str, payload: str, topic: str | None) -> None:
-    """Send a message via WebSocket."""
-    try:
-        if topic is None:
-            topic = determine_response_topic()
-
-        body = {"type": event_type, "payload": json.loads(payload), "topic": topic}
-        await ws_connection.send(json.dumps(body))
-    except WebSocketException as e:
-        logger.error("Error sending WebSocket message: %s", e)
-    except Exception as e:
-        logger.error("Unexpected error sending WebSocket message: %s", e)
-
-
-async def _send_subscribe_command(ws_connection: Any, topic: str) -> None:
-    """Send subscribe command via WebSocket."""
-    try:
-        body = {"type": "subscribe", "topic": topic, "payload": {}}
-        await ws_connection.send(json.dumps(body))
-        logger.debug("Subscribed to topic: %s", topic)
-    except WebSocketException as e:
-        logger.error("Error subscribing to topic %s: %s", topic, e)
-    except Exception as e:
-        logger.error("Unexpected error subscribing to topic %s: %s", topic, e)
-
-
-async def _send_unsubscribe_command(ws_connection: Any, topic: str) -> None:
-    """Send unsubscribe command via WebSocket."""
-    try:
-        body = {"type": "unsubscribe", "topic": topic, "payload": {}}
-        await ws_connection.send(json.dumps(body))
-        logger.debug("Unsubscribed from topic: %s", topic)
-    except WebSocketException as e:
-        logger.error("Error unsubscribing from topic %s: %s", topic, e)
-    except Exception as e:
-        logger.error("Unexpected error unsubscribing from topic %s: %s", topic, e)
-
-
 async def _process_event_queue() -> None:
     """Process events concurrently - runs on main thread."""
     logger.debug("Starting event queue processor on main thread")
@@ -382,13 +300,7 @@ async def _process_event_request(event: EventRequest) -> None:
         event.request,
         result_context={"response_topic": event.response_topic, "request_id": event.request_id},
     )
-
-    if result_event.result.succeeded():
-        dest_socket = "success_result"
-    else:
-        dest_socket = "failure_result"
-
-    await _send_message(dest_socket, result_event.json(), topic=result_event.response_topic)
+    await _process_node_event(GriptapeNodeEvent(wrapped_event=result_event))
 
 
 async def _process_app_event(event: AppEvent) -> None:
@@ -403,8 +315,21 @@ async def _process_node_event(event: GriptapeNodeEvent) -> None:
     """Process GriptapeNodeEvents and send them to the API (async version)."""
     # Emit the result back to the GUI
     result_event = event.wrapped_event
+
     if isinstance(result_event, EventResultSuccess):
         dest_socket = "success_result"
+        # Handle session-specific topic subscriptions
+        if isinstance(result_event.result, app_events.AppStartSessionResultSuccess):
+            session_id = result_event.result.session_id
+            topic = f"sessions/{session_id}/request"
+            await _subscribe_to_topic(topic)
+            logger.info("Subscribed to session topic: %s", topic)
+        elif isinstance(result_event.result, app_events.AppEndSessionResultSuccess):
+            session_id = result_event.result.session_id
+            if session_id is not None:
+                topic = f"sessions/{session_id}/request"
+                await _unsubscribe_from_topic(topic)
+                logger.info("Unsubscribed from session topic: %s", topic)
     elif isinstance(result_event, EventResultFailure):
         dest_socket = "failure_result"
     else:
@@ -443,14 +368,14 @@ async def _send_message(event_type: str, payload: str, topic: str | None = None)
 
     # Determine topic based on session_id and engine_id in the payload
     if topic is None:
-        topic = determine_response_topic()
+        topic = _determine_response_topic()
 
     message = WebSocketMessage(event_type, payload, topic)
 
     asyncio.run_coroutine_threadsafe(ws_outgoing_queue.put(message), websocket_event_loop)
 
 
-async def subscribe_to_topic(topic: str) -> None:
+async def _subscribe_to_topic(topic: str) -> None:
     """Queue a subscribe command for WebSocket using run_coroutine_threadsafe."""
     # Wait for websocket event loop to be ready
     websocket_event_loop_ready.wait()
@@ -462,7 +387,7 @@ async def subscribe_to_topic(topic: str) -> None:
     asyncio.run_coroutine_threadsafe(ws_outgoing_queue.put(SubscribeCommand(topic)), websocket_event_loop)
 
 
-async def unsubscribe_from_topic(topic: str) -> None:
+async def _unsubscribe_from_topic(topic: str) -> None:
     """Queue an unsubscribe command for WebSocket using run_coroutine_threadsafe."""
     if websocket_event_loop is None:
         logger.error("WebSocket event loop not available for unsubscribe")
@@ -471,7 +396,7 @@ async def unsubscribe_from_topic(topic: str) -> None:
     asyncio.run_coroutine_threadsafe(ws_outgoing_queue.put(UnsubscribeCommand(topic)), websocket_event_loop)
 
 
-def determine_response_topic() -> str | None:
+def _determine_response_topic() -> str:
     """Determine the response topic based on session_id and engine_id in the payload."""
     engine_id = griptape_nodes.get_engine_id()
     session_id = griptape_nodes.get_session_id()
@@ -487,21 +412,3 @@ def determine_response_topic() -> str | None:
 
     # Default to generic response topic
     return "response"
-
-
-def determine_request_topic() -> str | None:
-    """Determine the request topic based on session_id and engine_id in the payload."""
-    engine_id = griptape_nodes.get_engine_id()
-    session_id = griptape_nodes.get_session_id()
-
-    # Normal topic determination logic
-    # Check for session_id first (highest priority)
-    if session_id:
-        return f"sessions/{session_id}/request"
-
-    # Check for engine_id if no session_id
-    if engine_id:
-        return f"engines/{engine_id}/request"
-
-    # Default to generic request topic
-    return "request"

griptape_nodes/bootstrap/utils/python_subprocess_executor.py

@@ -69,7 +69,7 @@ class PythonSubprocessExecutor:
                 raise RuntimeError(msg)  # noqa: TRY301
 
         except Exception as e:
-            msg = "Error running subprocess"
+            msg = f"Error running subprocess: {e}"
             logger.exception(msg)
             raise PythonSubprocessExecutorError(msg) from e
         finally:

griptape_nodes/bootstrap/workflow_executors/local_session_workflow_executor.py

@@ -1,17 +1,13 @@
 from __future__ import annotations
 
 import asyncio
+import json
 import logging
 import threading
 from typing import TYPE_CHECKING, Any, Self
 
-from websockets.exceptions import ConnectionClosed, ConnectionClosedError
-
-from griptape_nodes.app.app import (
-    WebSocketMessage,
-    _create_websocket_connection,
-    _send_websocket_message,
-)
+from griptape_nodes.api_client import Client
+from griptape_nodes.app.app import WebSocketMessage
 from griptape_nodes.bootstrap.workflow_executors.local_workflow_executor import (
     LocalExecutorError,
     LocalWorkflowExecutor,
@@ -63,7 +59,7 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
 
         logger.info("Setting up session %s", self._session_id)
         GriptapeNodes.SessionManager().save_session(self._session_id)
-        GriptapeNodes.SessionManager().set_active_session_id(self._session_id)
+        GriptapeNodes.SessionManager().active_session_id = self._session_id
         await self._start_websocket_connection()
 
         return self
@@ -169,7 +165,10 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
         )
 
         # Send the run command to actually execute it (fire and forget)
-        start_flow_request = StartFlowRequest(flow_name=flow_name)
+        pickle_control_flow_result = kwargs.get("pickle_control_flow_result", False)
+        start_flow_request = StartFlowRequest(
+            flow_name=flow_name, pickle_control_flow_result=pickle_control_flow_result
+        )
         start_flow_task = asyncio.create_task(GriptapeNodes.ahandle_request(start_flow_request))
 
         is_flow_finished = False
@@ -261,13 +260,7 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
     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)
-        api_key = self._get_api_key()
-        if api_key is None:
-            logger.warning("No API key found, websocket connection will not be established")
-            return
-
-        logger.info("API key found, starting websocket thread")
-        self._websocket_thread = threading.Thread(target=self._start_websocket_thread, args=(api_key,), daemon=True)
+        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):
@@ -276,16 +269,7 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
         else:
             logger.error("Timeout waiting for websocket thread to start")
 
-    def _get_api_key(self) -> str | None:
-        """Get API key from secrets manager."""
-        try:
-            secrets_manager = GriptapeNodes.SecretsManager()
-            return secrets_manager.get_secret("GT_CLOUD_API_KEY")
-        except Exception:
-            logger.exception("Failed to get API key")
-            return None
-
-    def _start_websocket_thread(self, api_key: str) -> None:
+    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
@@ -302,7 +286,7 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
             logger.info("Websocket thread started and ready")
 
             # Run the async WebSocket tasks
-            loop.run_until_complete(self._run_websocket_tasks(api_key))
+            loop.run_until_complete(self._run_websocket_tasks())
         except Exception as e:
             logger.error("WebSocket thread error: %s", e)
         finally:
@@ -311,35 +295,21 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
             self._shutdown_event = None
             logger.info("Websocket thread ended")
 
-    async def _run_websocket_tasks(self, api_key: str) -> None:
+    async def _run_websocket_tasks(self) -> None:
         """Run websocket tasks - establish connection and handle outgoing messages."""
-        logger.info("Creating websocket connection stream")
-        connection_stream = _create_websocket_connection(api_key)
+        logger.info("Creating Client for session %s", self._session_id)
 
-        async for ws_connection in connection_stream:
+        async with Client() as client:
             logger.info("WebSocket connection established for session %s", self._session_id)
-            try:
-                # Use our own version that works with our local queue
-                await self._send_outgoing_messages(ws_connection)
 
-            except (ConnectionClosed, ConnectionClosedError):
-                logger.info("WebSocket connection closed, reconnecting...")
-                continue
-            except asyncio.CancelledError:
-                logger.info("WebSocket task cancelled, shutting down")
-                break
+            try:
+                await self._send_outgoing_messages(client)
             except Exception:
                 logger.exception("WebSocket tasks failed")
-                await asyncio.sleep(2.0)
-                continue
-
-            if self._shutdown_event and self._shutdown_event.is_set():
-                logger.info("Shutdown requested, ending WebSocket connection loop")
-                break
-
-        logger.info("WebSocket connection loop ended")
+            finally:
+                logger.info("WebSocket connection loop ended")
 
-    async def _send_outgoing_messages(self, ws_connection: Any) -> None:
+    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")
@@ -362,7 +332,9 @@ class LocalSessionWorkflowExecutor(LocalWorkflowExecutor):
 
             try:
                 if isinstance(message, WebSocketMessage):
-                    await _send_websocket_message(ws_connection, message.event_type, message.payload, message.topic)
+                    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))

griptape_nodes/bootstrap/workflow_executors/local_workflow_executor.py

@@ -104,6 +104,8 @@ class LocalWorkflowExecutor(WorkflowExecutor):
         for node_name, node in nodes.items():
             if isinstance(node, EndNode):
                 output[node_name] = node.parameter_values
+                # Parameter_output_values should also be included, and should take priority over parameter_values
+                output[node_name].update(node.parameter_output_values)
 
         return output
 
@@ -216,7 +218,10 @@ class LocalWorkflowExecutor(WorkflowExecutor):
         )
 
         # Now send the run command to actually execute it
-        start_flow_request = StartFlowRequest(flow_name=flow_name)
+        pickle_control_flow_result = kwargs.get("pickle_control_flow_result", False)
+        start_flow_request = StartFlowRequest(
+            flow_name=flow_name, pickle_control_flow_result=pickle_control_flow_result
+        )
         start_flow_result = await GriptapeNodes.ahandle_request(start_flow_request)
 
         if start_flow_result.failed():

griptape_nodes/bootstrap/workflow_executors/subprocess_workflow_executor.py

@@ -10,9 +10,8 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Self
 
 import anyio
-from websockets.exceptions import ConnectionClosed, ConnectionClosedError
 
-from griptape_nodes.app.app import _create_websocket_connection, _send_subscribe_command
+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.drivers.storage import StorageBackend
@@ -80,6 +79,8 @@ class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessE
         workflow_name: str,  # noqa: ARG002
         flow_input: Any,
         storage_backend: StorageBackend = StorageBackend.LOCAL,
+        *,
+        pickle_control_flow_result: bool = False,
         **kwargs: Any,  # noqa: ARG002
     ) -> None:
         """Execute a workflow in a subprocess and wait for completion."""
@@ -117,6 +118,9 @@ class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessE
                 str(tmp_workflow_path),
             ]
 
+            if pickle_control_flow_result:
+                args.append("--pickle-control-flow-result")
+
             try:
                 await self.execute_python_script(
                     script_path=tmp_script_path,
@@ -135,13 +139,7 @@ class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessE
     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)
-        api_key = self._get_api_key()
-        if api_key is None:
-            logger.warning("No API key found, WebSocket listener will not be started")
-            return
-
-        logger.info("API key found, starting WebSocket listener thread")
-        self._websocket_thread = threading.Thread(target=self._start_websocket_thread, args=(api_key,), daemon=True)
+        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):
@@ -173,7 +171,7 @@ class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessE
         else:
             logger.info("WebSocket listener thread stopped successfully")
 
-    def _start_websocket_thread(self, api_key: str) -> None:
+    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
@@ -189,7 +187,7 @@ class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessE
             logger.info("WebSocket listener thread started and ready")
 
             # Run the async WebSocket listener
-            loop.run_until_complete(self._run_websocket_listener(api_key))
+            loop.run_until_complete(self._run_websocket_listener())
         except Exception as e:
             logger.error("WebSocket listener thread error: %s", e)
         finally:
@@ -198,59 +196,37 @@ class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessE
             self._shutdown_event = None
             logger.info("WebSocket listener thread ended")
 
-    async def _run_websocket_listener(self, api_key: str) -> None:
+    async def _run_websocket_listener(self) -> None:
         """Run WebSocket listener - establish connection and handle incoming messages."""
-        logger.info("Creating WebSocket connection stream for listening")
-        connection_stream = _create_websocket_connection(api_key)
+        logger.info("Creating Client for listening on session %s", self._session_id)
 
-        async for ws_connection in connection_stream:
+        async with Client() as client:
             logger.info("WebSocket connection established for session %s", self._session_id)
+
             try:
-                # Listen for incoming messages
-                await self._listen_for_messages(ws_connection)
-
-            except (ConnectionClosed, ConnectionClosedError):
-                logger.info("WebSocket connection closed, reconnecting...")
-                continue
-            except asyncio.CancelledError:
-                logger.info("WebSocket listener task cancelled, shutting down")
-                break
+                await self._listen_for_messages(client)
             except Exception:
                 logger.exception("WebSocket listener failed")
-                await asyncio.sleep(2.0)
-                continue
-
-            # Check if shutdown was requested
-            if self._shutdown_event and self._shutdown_event.is_set():
-                logger.info("Shutdown requested, ending WebSocket listener connection loop")
-                break
-
-        logger.info("WebSocket listener connection loop ended")
+            finally:
+                logger.info("WebSocket listener connection loop ended")
 
-    async def _listen_for_messages(self, ws_connection: Any) -> None:
+    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")
 
-        # Subscribe to the session topic to receive messages
         topic = f"sessions/{self._session_id}/response"
+        await client.subscribe(topic)
 
         try:
-            await _send_subscribe_command(
-                ws_connection=ws_connection,
-                topic=topic,
-            )
-            async for message in ws_connection:
+            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:
-                    data = json.loads(message)
-                    logger.debug("Received WebSocket message: %s", data.get("type"))
-                    await self._process_event(data)
+                    logger.debug("Received WebSocket message: %s", message.get("type"))
+                    await self._process_event(message)
 
-                except json.JSONDecodeError:
-                    logger.warning("Failed to parse WebSocket message: %s", message)
                 except Exception:
                     logger.exception("Error processing WebSocket message")
 
@@ -285,7 +261,12 @@ class SubprocessWorkflowExecutor(LocalSessionWorkflowExecutor, PythonSubprocessE
 
         if isinstance(ex_event.payload, ControlFlowResolvedEvent):
             logger.info("Workflow execution completed successfully")
-            self.output = {ex_event.payload.end_node_name: ex_event.payload.parameter_output_values}
+            # Store both parameter output values and unique UUID values for deserialization
+            result = {
+                "parameter_output_values": ex_event.payload.parameter_output_values,
+                "unique_parameter_uuid_to_values": ex_event.payload.unique_parameter_uuid_to_values,
+            }
+            self.output = {ex_event.payload.end_node_name: result}
 
         if isinstance(ex_event.payload, ControlFlowCancelledEvent):
             logger.error("Workflow execution cancelled")