griptape-nodes 0.52.1__py3-none-any.whl → 0.54.0__py3-none-any.whl

griptape_nodes/machines/control_flow.py

@@ -9,7 +9,8 @@ from griptape_nodes.exe_types.core_types import Parameter
 from griptape_nodes.exe_types.node_types import BaseNode, NodeResolutionState
 from griptape_nodes.exe_types.type_validator import TypeValidator
 from griptape_nodes.machines.fsm import FSM, State
-from griptape_nodes.machines.node_resolution import NodeResolutionMachine
+from griptape_nodes.machines.parallel_resolution import ParallelResolutionMachine
+from griptape_nodes.machines.sequential_resolution import SequentialResolutionMachine
 from griptape_nodes.retained_mode.events.base_events import ExecutionEvent, ExecutionGriptapeNodeEvent
 from griptape_nodes.retained_mode.events.execution_events import (
     ControlFlowResolvedEvent,
@@ -17,6 +18,7 @@ from griptape_nodes.retained_mode.events.execution_events import (
     SelectedControlOutputEvent,
 )
 from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+from griptape_nodes.retained_mode.managers.settings import WorkflowExecutionMode
 
 
 @dataclass
@@ -37,40 +39,73 @@ logger = logging.getLogger("griptape_nodes")
 # This is the control flow context. Owns the Resolution Machine
 class ControlFlowContext:
     flow: ControlFlow
-    current_node: BaseNode | None
-    resolution_machine: NodeResolutionMachine
+    current_nodes: list[BaseNode]
+    resolution_machine: ParallelResolutionMachine | SequentialResolutionMachine
     selected_output: Parameter | None
     paused: bool = False
+    flow_name: str
+
+    def __init__(
+        self,
+        flow_name: str,
+        max_nodes_in_parallel: int,
+        *,
+        execution_type: WorkflowExecutionMode = WorkflowExecutionMode.SEQUENTIAL,
+    ) -> None:
+        self.flow_name = flow_name
+        if execution_type == WorkflowExecutionMode.PARALLEL:
+            # Get the global DagBuilder from FlowManager
+            from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+            dag_builder = GriptapeNodes.FlowManager().global_dag_builder
+            self.resolution_machine = ParallelResolutionMachine(
+                flow_name, max_nodes_in_parallel, dag_builder=dag_builder
+            )
+        else:
+            self.resolution_machine = SequentialResolutionMachine()
+        self.current_nodes = []
 
-    def __init__(self) -> None:
-        self.resolution_machine = NodeResolutionMachine()
-        self.current_node = None
-
-    def get_next_node(self, output_parameter: Parameter) -> NextNodeInfo | None:
-        """Get the next node and the target parameter that will receive the control flow.
+    def get_next_nodes(self, output_parameter: Parameter | None = None) -> list[NextNodeInfo]:
+        """Get all next nodes from the current nodes.
 
         Returns:
-            NextNodeInfo | None: Information about the next node or None if no connection
+            list[NextNodeInfo]: List of next nodes to process
         """
-        if self.current_node is not None:
-            node_connection = (
-                GriptapeNodes.FlowManager().get_connections().get_connected_node(self.current_node, output_parameter)
-            )
-            if node_connection is not None:
-                node, entry_parameter = node_connection
-                return NextNodeInfo(node=node, entry_parameter=entry_parameter)
-            # Continue Execution to the next node that needs to be executed using global execution queue
-            # Get the next node in the execution queue, or None if queue is empty
+        next_nodes = []
+        for current_node in self.current_nodes:
+            if output_parameter is not None:
+                # Get connected node from control flow
+                node_connection = (
+                    GriptapeNodes.FlowManager().get_connections().get_connected_node(current_node, output_parameter)
+                )
+                if node_connection is not None:
+                    node, entry_parameter = node_connection
+                    next_nodes.append(NextNodeInfo(node=node, entry_parameter=entry_parameter))
+            else:
+                # Get next control output for this node
+                next_output = current_node.get_next_control_output()
+                if next_output is not None:
+                    node_connection = (
+                        GriptapeNodes.FlowManager().get_connections().get_connected_node(current_node, next_output)
+                    )
+                    if node_connection is not None:
+                        node, entry_parameter = node_connection
+                        next_nodes.append(NextNodeInfo(node=node, entry_parameter=entry_parameter))
+
+        # If no connections found, check execution queue
+        if not next_nodes:
             node = GriptapeNodes.FlowManager().get_next_node_from_execution_queue()
             if node is not None:
-                return NextNodeInfo(node=node, entry_parameter=None)
-        return None
+                next_nodes.append(NextNodeInfo(node=node, entry_parameter=None))
+
+        return next_nodes
 
-    def reset(self) -> None:
-        if self.current_node:
-            self.current_node.clear_node()
-        self.current_node = None
-        self.resolution_machine.reset_machine()
+    def reset(self, *, cancel: bool = False) -> None:
+        if self.current_nodes is not None:
+            for node in self.current_nodes:
+                node.clear_node()
+        self.current_nodes = []
+        self.resolution_machine.reset_machine(cancel=cancel)
         self.selected_output = None
         self.paused = False
 
@@ -80,24 +115,25 @@ class ResolveNodeState(State):
     @staticmethod
     async def on_enter(context: ControlFlowContext) -> type[State] | None:
         # The state machine has started, but it hasn't began to execute yet.
-        if context.current_node is None:
+        if len(context.current_nodes) == 0:
             # We don't have anything else to do. Move back to Complete State so it has to restart.
             return CompleteState
 
-        # Mark the node unresolved, and broadcast an event to the GUI.
-        if not context.current_node.lock:
-            context.current_node.make_node_unresolved(
-                current_states_to_trigger_change_event=set(
-                    {NodeResolutionState.UNRESOLVED, NodeResolutionState.RESOLVED, NodeResolutionState.RESOLVING}
+        # Mark all current nodes unresolved and broadcast events
+        for current_node in context.current_nodes:
+            if not current_node.lock:
+                current_node.make_node_unresolved(
+                    current_states_to_trigger_change_event=set(
+                        {NodeResolutionState.UNRESOLVED, NodeResolutionState.RESOLVED, NodeResolutionState.RESOLVING}
+                    )
+                )
+            # Now broadcast that we have a current control node.
+            GriptapeNodes.EventManager().put_event(
+                ExecutionGriptapeNodeEvent(
+                    wrapped_event=ExecutionEvent(payload=CurrentControlNodeEvent(node_name=current_node.name))
                 )
             )
-        # Now broadcast that we have a current control node.
-        GriptapeNodes.EventManager().put_event(
-            ExecutionGriptapeNodeEvent(
-                wrapped_event=ExecutionEvent(payload=CurrentControlNodeEvent(node_name=context.current_node.name))
-            )
-        )
-        logger.info("Resolving %s", context.current_node.name)
+            logger.info("Resolving %s", current_node.name)
         if not context.paused:
             # Call the update. Otherwise wait
             return ResolveNodeState
@@ -106,13 +142,17 @@ class ResolveNodeState(State):
     # This is necessary to transition to the next step.
     @staticmethod
     async def on_update(context: ControlFlowContext) -> type[State] | None:
-        # If node has not already been resolved!
-        if context.current_node is None:
+        # If no current nodes, we're done
+        if len(context.current_nodes) == 0:
             return CompleteState
-        if context.current_node.state != NodeResolutionState.RESOLVED:
-            await context.resolution_machine.resolve_node(context.current_node)
+
+        # Resolve nodes - pass first node for sequential resolution
+        current_node = context.current_nodes[0] if context.current_nodes else None
+        await context.resolution_machine.resolve_node(current_node)
 
         if context.resolution_machine.is_complete():
+            if isinstance(context.resolution_machine, ParallelResolutionMachine):
+                return CompleteState
             return NextNodeState
         return None
 
@@ -120,44 +160,49 @@ class ResolveNodeState(State):
 class NextNodeState(State):
     @staticmethod
     async def on_enter(context: ControlFlowContext) -> type[State] | None:
-        if context.current_node is None:
+        if len(context.current_nodes) == 0:
             return CompleteState
-        # I did define this on the ControlNode.
-        if context.current_node.stop_flow:
-            # We're done here.
-            context.current_node.stop_flow = False
+
+        # Check for stop_flow on any current nodes
+        for current_node in context.current_nodes[:]:
+            if current_node.stop_flow:
+                current_node.stop_flow = False
+                context.current_nodes.remove(current_node)
+
+        # If all nodes stopped flow, complete
+        if len(context.current_nodes) == 0:
             return CompleteState
-        next_output = context.current_node.get_next_control_output()
-        next_node_info = None
 
-        if next_output is not None:
-            context.selected_output = next_output
-            next_node_info = context.get_next_node(context.selected_output)
-            GriptapeNodes.EventManager().put_event(
-                ExecutionGriptapeNodeEvent(
-                    wrapped_event=ExecutionEvent(
-                        payload=SelectedControlOutputEvent(
-                            node_name=context.current_node.name,
-                            selected_output_parameter_name=next_output.name,
+        # Get all next nodes from current nodes
+        next_node_infos = context.get_next_nodes()
+
+        # Broadcast selected control output events for nodes with outputs
+        for current_node in context.current_nodes:
+            next_output = current_node.get_next_control_output()
+            if next_output is not None:
+                context.selected_output = next_output
+                GriptapeNodes.EventManager().put_event(
+                    ExecutionGriptapeNodeEvent(
+                        wrapped_event=ExecutionEvent(
+                            payload=SelectedControlOutputEvent(
+                                node_name=current_node.name,
+                                selected_output_parameter_name=next_output.name,
+                            )
                         )
                     )
                 )
-            )
-        else:
-            # Get the next node in the execution queue, or None if queue is empty
-            next_node = GriptapeNodes.FlowManager().get_next_node_from_execution_queue()
-            if next_node is not None:
-                next_node_info = NextNodeInfo(node=next_node, entry_parameter=None)
-
-        # The parameter that will be evaluated next
-        if next_node_info is None:
-            # If no node attached
+
+        # If no next nodes, we're complete
+        if not next_node_infos:
             return CompleteState
 
-        # Always set the entry control parameter (None for execution queue nodes)
-        next_node_info.node.set_entry_control_parameter(next_node_info.entry_parameter)
+        # Set up next nodes as current nodes
+        next_nodes = []
+        for next_node_info in next_node_infos:
+            next_node_info.node.set_entry_control_parameter(next_node_info.entry_parameter)
+            next_nodes.append(next_node_info.node)
 
-        context.current_node = next_node_info.node
+        context.current_nodes = next_nodes
         context.selected_output = None
         if not context.paused:
             return ResolveNodeState
@@ -171,15 +216,14 @@ class NextNodeState(State):
 class CompleteState(State):
     @staticmethod
     async def on_enter(context: ControlFlowContext) -> type[State] | None:
-        if context.current_node is not None:
+        # Broadcast completion events for any remaining current nodes
+        for current_node in context.current_nodes:
             GriptapeNodes.EventManager().put_event(
                 ExecutionGriptapeNodeEvent(
                     wrapped_event=ExecutionEvent(
                         payload=ControlFlowResolvedEvent(
-                            end_node_name=context.current_node.name,
-                            parameter_output_values=TypeValidator.safe_serialize(
-                                context.current_node.parameter_output_values
-                            ),
+                            end_node_name=current_node.name,
+                            parameter_output_values=TypeValidator.safe_serialize(current_node.parameter_output_values),
                         )
                     )
                 )
@@ -194,14 +238,24 @@ class CompleteState(State):
 
 # MACHINE TIME!!!
 class ControlFlowMachine(FSM[ControlFlowContext]):
-    def __init__(self) -> None:
-        context = ControlFlowContext()
+    def __init__(self, flow_name: str) -> None:
+        execution_type = GriptapeNodes.ConfigManager().get_config_value(
+            "workflow_execution_mode", default=WorkflowExecutionMode.SEQUENTIAL
+        )
+        max_nodes_in_parallel = GriptapeNodes.ConfigManager().get_config_value("max_nodes_in_parallel", default=5)
+        context = ControlFlowContext(flow_name, max_nodes_in_parallel, execution_type=execution_type)
         super().__init__(context)
 
     async def start_flow(self, start_node: BaseNode, debug_mode: bool = False) -> None:  # noqa: FBT001, FBT002
-        self._context.current_node = start_node
+        # If using DAG resolution, process data_nodes from queue first
+        if isinstance(self._context.resolution_machine, ParallelResolutionMachine):
+            current_nodes = await self._process_nodes_for_dag(start_node)
+        else:
+            current_nodes = [start_node]
+        self._context.current_nodes = current_nodes
         # Set entry control parameter for initial node (None for workflow start)
-        start_node.set_entry_control_parameter(None)
+        for node in current_nodes:
+            node.set_entry_control_parameter(None)
         # Set up to debug
         self._context.paused = debug_mode
         await self.start(ResolveNodeState)  # Begins the flow
@@ -214,31 +268,78 @@ class ControlFlowMachine(FSM[ControlFlowContext]):
 
     def change_debug_mode(self, debug_mode: bool) -> None:  # noqa: FBT001
         self._context.paused = debug_mode
-        self._context.resolution_machine.change_debug_mode(debug_mode)
+        self._context.resolution_machine.change_debug_mode(debug_mode=debug_mode)
 
     async def granular_step(self, change_debug_mode: bool) -> None:  # noqa: FBT001
         resolution_machine = self._context.resolution_machine
+
         if change_debug_mode:
-            resolution_machine.change_debug_mode(True)
+            resolution_machine.change_debug_mode(debug_mode=True)
         await resolution_machine.update()
 
-        # Tick the control flow if the resolution machine inside it isn't busy.
-        if resolution_machine.is_complete() or not resolution_machine.is_started():  # noqa: SIM102
+        # Tick the control flow if the current machine isn't busy
+        if self._current_state is ResolveNodeState and (  # noqa: SIM102
+            resolution_machine.is_complete() or not resolution_machine.is_started()
+        ):
             # Don't tick ourselves if we are already complete.
             if self._current_state is not None:
                 await self.update()
 
     async def node_step(self) -> None:
         resolution_machine = self._context.resolution_machine
-        resolution_machine.change_debug_mode(False)
-        await resolution_machine.update()
 
-        # Tick the control flow if the resolution machine inside it isn't busy.
-        if resolution_machine.is_complete() or not resolution_machine.is_started():  # noqa: SIM102
-            # Don't tick ourselves if we are already complete.
-            if self._current_state is not None:
-                await self.update()
+        resolution_machine.change_debug_mode(debug_mode=False)
+
+        # If we're in the resolution phase, step the resolution machine
+        if self._current_state is ResolveNodeState:
+            await resolution_machine.update()
 
-    def reset_machine(self) -> None:
-        self._context.reset()
+        # Tick the control flow if the current machine isn't busy
+        if self._current_state is ResolveNodeState and (
+            resolution_machine.is_complete() or not resolution_machine.is_started()
+        ):
+            await self.update()
+
+    async def _process_nodes_for_dag(self, start_node: BaseNode) -> list[BaseNode]:
+        """Process data_nodes from the global queue to build unified DAG.
+
+        This method identifies data_nodes in the execution queue and processes
+        their dependencies into the DAG resolution machine.
+        """
+        if not isinstance(self._context.resolution_machine, ParallelResolutionMachine):
+            return []
+        # Get the global flow queue
+        flow_manager = GriptapeNodes.FlowManager()
+        dag_builder = flow_manager.global_dag_builder
+        if dag_builder is None:
+            msg = "DAG builder is not initialized."
+            raise ValueError(msg)
+        # Build with the first node:
+        dag_builder.add_node_with_dependencies(start_node, start_node.name)
+        queue_items = list(flow_manager.global_flow_queue.queue)
+        start_nodes = [start_node]
+        # Find data_nodes and remove them from queue
+        for item in queue_items:
+            from griptape_nodes.retained_mode.managers.flow_manager import DagExecutionType
+
+            if item.dag_execution_type in (DagExecutionType.CONTROL_NODE, DagExecutionType.START_NODE):
+                node = item.node
+                node.state = NodeResolutionState.UNRESOLVED
+                dag_builder.add_node_with_dependencies(node, node.name)
+                flow_manager.global_flow_queue.queue.remove(item)
+                start_nodes.append(node)
+            elif item.dag_execution_type == DagExecutionType.DATA_NODE:
+                node = item.node
+                node.state = NodeResolutionState.UNRESOLVED
+                # Build here.
+                dag_builder.add_node_with_dependencies(node, node.name)
+                flow_manager.global_flow_queue.queue.remove(item)
+        return start_nodes
+
+    def reset_machine(self, *, cancel: bool = False) -> None:
+        self._context.reset(cancel=cancel)
         self._current_state = None
+
+    @property
+    def resolution_machine(self) -> ParallelResolutionMachine | SequentialResolutionMachine:
+        return self._context.resolution_machine

griptape_nodes/machines/dag_builder.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import TYPE_CHECKING
+
+from griptape_nodes.common.directed_graph import DirectedGraph
+from griptape_nodes.exe_types.core_types import ParameterTypeBuiltin
+from griptape_nodes.exe_types.node_types import NodeResolutionState
+
+if TYPE_CHECKING:
+    import asyncio
+
+    from griptape_nodes.exe_types.connections import Connections
+    from griptape_nodes.exe_types.node_types import BaseNode
+
+logger = logging.getLogger("griptape_nodes")
+
+
+class NodeState(StrEnum):
+    """Individual node execution states."""
+
+    QUEUED = "queued"
+    PROCESSING = "processing"
+    DONE = "done"
+    CANCELED = "canceled"
+    ERRORED = "errored"
+    WAITING = "waiting"
+
+
+@dataclass(kw_only=True)
+class DagNode:
+    """Represents a node in the DAG with runtime references."""
+
+    task_reference: asyncio.Task | None = field(default=None)
+    node_state: NodeState = field(default=NodeState.WAITING)
+    node_reference: BaseNode
+
+
+class DagBuilder:
+    """Handles DAG construction independently of execution state machine."""
+
+    graphs: dict[str, DirectedGraph]  # Str is the name of the start node associated here.
+    node_to_reference: dict[str, DagNode]
+
+    def __init__(self) -> None:
+        self.graphs = {}
+        self.node_to_reference: dict[str, DagNode] = {}
+
+    # Complex with the inner recursive method, but it needs connections and added_nodes.
+    def add_node_with_dependencies(self, node: BaseNode, graph_name: str = "default") -> list[BaseNode]:  # noqa: C901
+        """Add node and all its dependencies to DAG. Returns list of added nodes."""
+        from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+        connections = GriptapeNodes.FlowManager().get_connections()
+        added_nodes = []
+        graph = self.graphs.get(graph_name, None)
+        if graph is None:
+            graph = DirectedGraph()
+            self.graphs[graph_name] = graph
+
+        def _add_node_recursive(current_node: BaseNode, visited: set[str], graph: DirectedGraph) -> None:
+            if current_node.name in visited:
+                return
+            visited.add(current_node.name)
+
+            # Skip if already in DAG (use DAG membership, not resolved state)
+            if current_node.name in self.node_to_reference:
+                return
+
+            # Process dependencies first (depth-first)
+            ignore_data_dependencies = False
+            # This is specifically for output_selector. Overriding 'initialize_spotlight' doesn't work anymore.
+            if hasattr(current_node, "ignore_dependencies"):
+                ignore_data_dependencies = True
+            for param in current_node.parameters:
+                if param.type == ParameterTypeBuiltin.CONTROL_TYPE:
+                    continue
+                if ignore_data_dependencies:
+                    continue
+                upstream_connection = connections.get_connected_node(current_node, param)
+                if upstream_connection:
+                    upstream_node, _ = upstream_connection
+                    # Don't add nodes that have already been resolved.
+                    if upstream_node.state == NodeResolutionState.RESOLVED:
+                        continue
+                    # If upstream is already in DAG, skip creating edge (it's in another graph)
+                    if upstream_node.name in self.node_to_reference:
+                        graph.add_edge(upstream_node.name, current_node.name)
+                    # Otherwise, add it to DAG first then create edge
+                    else:
+                        _add_node_recursive(upstream_node, visited, graph)
+                        graph.add_edge(upstream_node.name, current_node.name)
+
+            # Add current node to DAG (but keep original resolution state)
+
+            dag_node = DagNode(node_reference=current_node, node_state=NodeState.WAITING)
+            self.node_to_reference[current_node.name] = dag_node
+            graph.add_node(node_for_adding=current_node.name)
+            # DON'T mark as resolved - that happens during actual execution
+            added_nodes.append(current_node)
+
+        _add_node_recursive(node, set(), graph)
+
+        return added_nodes
+
+    def add_node(self, node: BaseNode, graph_name: str = "default") -> DagNode:
+        """Add just one node to DAG without dependencies (assumes dependencies already exist)."""
+        if node.name in self.node_to_reference:
+            return self.node_to_reference[node.name]
+
+        dag_node = DagNode(node_reference=node, node_state=NodeState.WAITING)
+        self.node_to_reference[node.name] = dag_node
+        graph = self.graphs.get(graph_name, None)
+        if graph is None:
+            graph = DirectedGraph()
+            self.graphs[graph_name] = graph
+        graph.add_node(node_for_adding=node.name)
+        return dag_node
+
+    def clear(self) -> None:
+        """Clear all nodes and references from the DAG builder."""
+        self.graphs.clear()
+        self.node_to_reference.clear()
+
+    def can_queue_control_node(self, node: DagNode) -> bool:
+        if len(self.graphs) == 1:
+            return True
+
+        from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+        connections = GriptapeNodes.FlowManager().get_connections()
+
+        control_connections = self.get_number_incoming_control_connections(node.node_reference, connections)
+        if control_connections <= 1:
+            return True
+
+        for graph in self.graphs.values():
+            # If the length of the graph is 0, skip it. it's either reached it or it's a dead end.
+            if len(graph.nodes()) == 0:
+                continue
+
+            # If graph has nodes, the root node (not the leaf, the root), check forward path from that
+            root_nodes = [n for n in graph.nodes() if graph.out_degree(n) == 0]
+            for root_node_name in root_nodes:
+                if root_node_name in self.node_to_reference:
+                    root_node = self.node_to_reference[root_node_name].node_reference
+
+                    # Skip if the root node is the same as the target node - it can't reach itself
+                    if root_node == node.node_reference:
+                        continue
+
+                    # Check if the target node is in the forward path from this root
+                    if self._is_node_in_forward_path(root_node, node.node_reference, connections):
+                        return False  # This graph could still reach the target node
+
+        # Otherwise, return true at the end of the function
+        return True
+
+    def get_number_incoming_control_connections(self, node: BaseNode, connections: Connections) -> int:
+        if node.name not in connections.incoming_index:
+            return 0
+
+        control_connection_count = 0
+        node_connections = connections.incoming_index[node.name]
+
+        for param_name, connection_ids in node_connections.items():
+            # Find the parameter to check if it's a control type
+            param = node.get_parameter_by_name(param_name)
+            if param and ParameterTypeBuiltin.CONTROL_TYPE.value in param.input_types:
+                control_connection_count += len(connection_ids)
+
+        return control_connection_count
+
+    def _is_node_in_forward_path(
+        self, start_node: BaseNode, target_node: BaseNode, connections: Connections, visited: set[str] | None = None
+    ) -> bool:
+        """Check if target_node is reachable from start_node through control flow connections."""
+        if visited is None:
+            visited = set()
+
+        if start_node.name in visited:
+            return False
+        visited.add(start_node.name)
+
+        # Check ALL outgoing control connections, not just get_next_control_output()
+        # This handles IfElse nodes that have multiple possible control outputs
+        if start_node.name in connections.outgoing_index:
+            for param_name, connection_ids in connections.outgoing_index[start_node.name].items():
+                # Find the parameter to check if it's a control type
+                param = start_node.get_parameter_by_name(param_name)
+                if param and param.output_type == ParameterTypeBuiltin.CONTROL_TYPE.value:
+                    # This is a control parameter - check all its connections
+                    for connection_id in connection_ids:
+                        if connection_id in connections.connections:
+                            connection = connections.connections[connection_id]
+                            next_node = connection.target_node
+
+                            if next_node.name == target_node.name:
+                                return True
+
+                            # Recursively check the forward path
+                            if self._is_node_in_forward_path(next_node, target_node, connections, visited):
+                                return True
+
+        return False

griptape_nodes/machines/fsm.py

@@ -34,9 +34,18 @@ class FSM[T]:
         # Enter the initial state.
         await self.transition_state(initial_state)
 
-    def get_current_state(self) -> type[State] | None:
+    @property
+    def current_state(self) -> type[State] | None:
         return self._current_state
 
+    @current_state.setter
+    def current_state(self, value: type[State] | None) -> None:
+        self._current_state = value
+
+    @property
+    def context(self) -> T:
+        return self._context
+
     async def transition_state(self, new_state: type[State] | None) -> None:
         while new_state is not None:
             # Exit the current state.