griptape-nodes 0.55.1__py3-none-any.whl → 0.56.1__py3-none-any.whl

griptape_nodes/cli/commands/init.py

@@ -55,6 +55,10 @@ def init_command(  # noqa: PLR0913
         bool,
         typer.Option(help="Run init in non-interactive mode (no prompts)."),
     ] = False,
+    hf_token: Annotated[
+        str | None,
+        typer.Option(help="Set the Hugging Face token for downloading gated models."),
+    ] = None,
     config: Annotated[
         list[str] | None,
         typer.Option(
@@ -83,6 +87,7 @@ def init_command(  # noqa: PLR0913
             secret_values=secret_values,
             libraries_sync=libraries_sync,
             bucket_name=bucket_name,
+            hf_token=hf_token,
         )
     )
 
@@ -128,6 +133,7 @@ def _run_init_configuration(config: InitConfig) -> None:
     _handle_workspace_config(config)
     _handle_storage_backend_config(config)
     _handle_bucket_config(config)
+    _handle_hf_token_config(config)
     _handle_advanced_library_config(config)
     _handle_arbitrary_configs(config)
 
@@ -193,6 +199,25 @@ def _handle_bucket_config(config: InitConfig) -> str | None:
     return bucket_id
 
 
+def _handle_hf_token_config(config: InitConfig) -> str | None:
+    """Handle Hugging Face token configuration step."""
+    hf_token = None
+
+    if config.interactive:
+        # First ask if they want to configure an HF token
+        configure_hf_token = _prompt_for_hf_token_configuration()
+        if configure_hf_token:
+            hf_token = _prompt_for_hf_token(default_hf_token=config.hf_token)
+    elif config.hf_token is not None:
+        hf_token = config.hf_token
+
+    if hf_token is not None:
+        secrets_manager.set_secret("HF_TOKEN", hf_token)
+        console.print("[bold green]Hugging Face token set[/bold green]")
+
+    return hf_token
+
+
 def _handle_advanced_library_config(config: InitConfig) -> bool | None:
     """Handle advanced library configuration step."""
     register_advanced_library = config.register_advanced_library
@@ -513,6 +538,69 @@ def _create_new_bucket(bucket_name: str) -> str:
         return bucket_id
 
 
+def _prompt_for_hf_token_configuration() -> bool:
+    """Prompts the user whether to configure a Hugging Face token."""
+    # Check if there's already an HF token configured
+    current_hf_token = secrets_manager.get_secret("HF_TOKEN", should_error_on_not_found=False)
+
+    if current_hf_token:
+        explainer = """[bold cyan]Hugging Face Token Configuration[/bold cyan]
+    You currently have a Hugging Face token configured.
+
+    Hugging Face tokens are used to access gated models from the Hugging Face Hub, such as:
+    - Meta's Llama models
+    - black-forest-labs/FLUX.1-dev
+    - Other restricted or premium models
+
+    Would you like to update your Hugging Face token or keep the current one?"""
+        prompt_text = "Update Hugging Face token?"
+        default_value = False
+    else:
+        explainer = """[bold cyan]Hugging Face Token Configuration[/bold cyan]
+    Would you like to configure a Hugging Face token?
+
+    Hugging Face tokens are used by the model manager to download gated models from the Hugging Face Hub, such as:
+    - Meta's Llama models
+    - black-forest-labs/FLUX.1-dev
+    - Other restricted or premium models
+
+    If you don't plan to use gated models, you can skip this step.
+    You can get a token from https://huggingface.co/settings/tokens
+
+    You can always configure a token later by running the initialization process again."""
+        prompt_text = "Configure Hugging Face token?"
+        default_value = False
+
+    console.print(Panel(explainer, expand=False))
+    return Confirm.ask(prompt_text, default=default_value)
+
+
+def _prompt_for_hf_token(default_hf_token: str | None = None) -> str | None:
+    """Prompts the user for their Hugging Face token."""
+    if default_hf_token is None:
+        default_hf_token = secrets_manager.get_secret("HF_TOKEN", should_error_on_not_found=False)
+
+    explainer = """[bold cyan]Hugging Face Token[/bold cyan]
+    Please enter your Hugging Face token to enable downloading of gated models.
+
+    To get a token:
+    1. Go to https://huggingface.co/settings/tokens
+    2. Create a new token with 'Read' permissions
+    3. Copy and paste the token here
+
+    You can leave this blank to skip token configuration."""
+    console.print(Panel(explainer, expand=False))
+
+    hf_token = Prompt.ask(
+        "Hugging Face Token (optional)",
+        default=default_hf_token or "",
+        show_default=False,
+    )
+
+    # Return None if empty string
+    return hf_token if hf_token.strip() else None
+
+
 def _parse_key_value_pairs(pairs: list[str] | None) -> dict[str, Any] | None:
     """Parse key=value pairs from a list of strings.
 

griptape_nodes/cli/commands/models.py

@@ -1,6 +1,7 @@
 """Models command for managing AI models."""
 
 import asyncio
+import sys
 from typing import TYPE_CHECKING
 
 import typer
@@ -130,6 +131,7 @@ async def _download_model(
     except Exception as e:
         console.print("[bold red]Model download failed:[/bold red]")
         console.print(f"[red]{e}[/red]")
+        sys.exit(1)
 
 
 async def _list_models() -> None:

griptape_nodes/cli/main.py

@@ -10,6 +10,7 @@ from rich.console import Console
 sys.path.append(str(Path.cwd()))
 
 from griptape_nodes.cli.commands import config, engine, init, libraries, models, self
+from griptape_nodes.cli.commands.engine import _auto_update_self
 from griptape_nodes.utils.version_utils import get_complete_version_string
 
 console = Console()
@@ -47,9 +48,13 @@ def main(
         console.print(f"[bold green]{version_string}[/bold green]")
         raise typer.Exit
 
+    # Run auto-update check for any command (unless disabled)
+    if not no_update:
+        _auto_update_self()
+
     if ctx.invoked_subcommand is None:
         # Default to engine command when no subcommand is specified
-        engine.engine_command(no_update=no_update)
+        engine.engine_command()
 
 
 if __name__ == "__main__":

griptape_nodes/cli/shared.py

@@ -22,6 +22,7 @@ class InitConfig:
     secret_values: dict[str, str] | None = None
     libraries_sync: bool | None = None
     bucket_name: str | None = None
+    hf_token: str | None = None
 
 
 # Initialize console

griptape_nodes/exe_types/core_types.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import uuid
 from abc import ABC, abstractmethod
 from copy import deepcopy
@@ -9,6 +10,8 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, Self, Type
 
 from pydantic import BaseModel
 
+logger = logging.getLogger("griptape_nodes")
+
 
 class NodeMessagePayload(BaseModel):
     """Structured payload for node messages.
@@ -554,6 +557,7 @@ class ParameterMessage(BaseNodeElement, UIOptionsMixin):
         button_align: ButtonAlignType = "full-width",
         full_width: bool = False,
         ui_options: dict | None = None,
+        traits: set[Trait.__class__ | Trait] | None = None,
         **kwargs,
     ):
         super().__init__(element_type=ParameterMessage.__name__, **kwargs)
@@ -569,6 +573,17 @@ class ParameterMessage(BaseNodeElement, UIOptionsMixin):
         self._full_width = full_width
         self._ui_options = ui_options or {}
 
+        # Handle traits if provided
+        if traits:
+            for trait in traits:
+                if isinstance(trait, type):
+                    # It's a trait class, instantiate it
+                    trait_instance = trait()
+                else:
+                    # It's already a trait instance
+                    trait_instance = trait
+                self.add_child(trait_instance)
+
     @property
     def variant(self) -> VariantType:
         return self._variant
@@ -694,6 +709,16 @@ class ParameterMessage(BaseNodeElement, UIOptionsMixin):
         else:
             button_icon = self.button_icon
 
+        # Check if there are any Button traits with on_click callbacks
+        has_button_callback = False
+        for child in self.children:
+            # Import here to avoid circular imports
+            from griptape_nodes.traits.button import Button
+
+            if isinstance(child, Button) and child.on_click_callback is not None:
+                has_button_callback = True
+                break
+
         # Merge the UI options with the message-specific options
         # Always include these fields, even if they're None or empty
         message_ui_options = {
@@ -705,6 +730,7 @@ class ParameterMessage(BaseNodeElement, UIOptionsMixin):
             "button_icon": button_icon,
             "button_variant": self.button_variant,
             "button_align": self.button_align,
+            "button_on_click": has_button_callback,
             "full_width": self.full_width,
         }
 
@@ -1483,12 +1509,23 @@ class ParameterList(ParameterContainer):
         user_defined: bool = False,
         element_id: str | None = None,
         element_type: str | None = None,
+        # UI convenience parameters
+        collapsed: bool | None = None,
+        child_prefix: str | None = None,
+        grid: bool | None = None,
+        grid_columns: int | None = None,
     ):
         if traits:
             self._original_traits = traits
         else:
             self._original_traits = set()
 
+        # Store the UI convenience parameters
+        self._collapsed = collapsed
+        self._child_prefix = child_prefix
+        self._grid = grid
+        self._grid_columns = grid_columns
+
         # Remember: we're a Parameter, too, just like everybody else.
         super().__init__(
             name=name,
@@ -1511,6 +1548,99 @@ class ParameterList(ParameterContainer):
             element_type=element_type,
         )
 
+    @property
+    def collapsed(self) -> bool | None:
+        return self._collapsed
+
+    @collapsed.setter
+    @BaseNodeElement.emits_update_on_write
+    def collapsed(self, value: bool | None) -> None:
+        self._collapsed = value
+
+    @property
+    def child_prefix(self) -> str | None:
+        return self._child_prefix
+
+    @child_prefix.setter
+    @BaseNodeElement.emits_update_on_write
+    def child_prefix(self, value: str | None) -> None:
+        self._child_prefix = value
+
+    @property
+    def grid(self) -> bool | None:
+        return self._grid
+
+    @grid.setter
+    @BaseNodeElement.emits_update_on_write
+    def grid(self, value: bool | None) -> None:
+        self._grid = value
+
+    @property
+    def grid_columns(self) -> int | None:
+        return self._grid_columns
+
+    @grid_columns.setter
+    @BaseNodeElement.emits_update_on_write
+    def grid_columns(self, value: int | None) -> None:
+        self._grid_columns = value
+
+    @property
+    def ui_options(self) -> dict:
+        """Override ui_options to merge convenience parameters in real-time."""
+        # Get base ui_options from parent
+        base_ui_options = super().ui_options
+
+        # Build convenience options from instance parameters
+        convenience_options = {}
+
+        if self._collapsed is not None:
+            convenience_options["collapsed"] = self._collapsed
+
+        if self._child_prefix is not None:
+            convenience_options["child_prefix"] = self._child_prefix
+
+        if self._grid is not None and self._grid:
+            convenience_options["display"] = "grid"
+
+        if self._grid_columns is not None and self._grid:
+            convenience_options["columns"] = self._grid_columns
+
+        # Merge convenience options with base ui_options
+        return {
+            **base_ui_options,
+            **convenience_options,
+        }
+
+    @ui_options.setter
+    @BaseNodeElement.emits_update_on_write
+    def ui_options(self, value: dict) -> None:
+        """Set ui_options, preserving convenience parameters."""
+        # Extract convenience parameters from the incoming value
+        if "display" in value and value["display"] == "grid":
+            self._grid = True
+            if "columns" in value:
+                self._grid_columns = value["columns"]
+        else:
+            self._grid = False
+
+        if "collapsed" in value:
+            self._collapsed = value["collapsed"]
+
+        if "child_prefix" in value:
+            self._child_prefix = value["child_prefix"]
+
+        # Set the base ui_options (excluding convenience parameters)
+        base_ui_options = {
+            k: v for k, v in value.items() if k not in ["display", "columns", "collapsed", "child_prefix"]
+        }
+        self._ui_options = base_ui_options
+
+    def to_dict(self) -> dict[str, Any]:
+        """Override to_dict to use the merged ui_options."""
+        data = super().to_dict()
+        data["ui_options"] = self.ui_options
+        return data
+
     def _custom_getter_for_property_type(self) -> str:
         base_type = super()._custom_getter_for_property_type()
         result = f"list[{base_type}]"

griptape_nodes/exe_types/node_types.py

@@ -5,8 +5,9 @@ import logging
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Generator, Iterable
 from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
 from enum import StrEnum, auto
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
 
 from griptape_nodes.exe_types.core_types import (
     BaseNodeElement,
@@ -43,6 +44,7 @@ from griptape_nodes.traits.options import Options
 
 if TYPE_CHECKING:
     from griptape_nodes.exe_types.core_types import NodeMessagePayload
+    from griptape_nodes.node_library.library_registry import LibraryNameAndVersion
 
 logger = logging.getLogger("griptape_nodes")
 
@@ -50,6 +52,53 @@ T = TypeVar("T")
 
 AsyncResult = Generator[Callable[[], T], T]
 
+LOCAL_EXECUTION = "Local Execution"
+
+
+class ImportDependency(NamedTuple):
+    """Import dependency specification for a node.
+
+    Attributes:
+        module: The module name to import
+        class_name: Optional class name to import from the module. If None, imports the entire module.
+    """
+
+    module: str
+    class_name: str | None = None
+
+
+@dataclass
+class NodeDependencies:
+    """Dependencies that a node has on external resources.
+
+    This class provides a way for nodes to declare their dependencies on workflows,
+    static files, Python imports, and libraries. This information can be used by the system
+    for workflow packaging, dependency resolution, and deployment planning.
+
+    Attributes:
+        referenced_workflows: Set of workflow names that this node references
+        static_files: Set of static file names that this node depends on
+        imports: Set of Python imports that this node requires
+        libraries: Set of library names and versions that this node uses
+    """
+
+    referenced_workflows: set[str] = field(default_factory=set)
+    static_files: set[str] = field(default_factory=set)
+    imports: set[ImportDependency] = field(default_factory=set)
+    libraries: set[LibraryNameAndVersion] = field(default_factory=set)
+
+    def aggregate_from(self, other: NodeDependencies) -> None:
+        """Aggregate dependencies from another NodeDependencies object into this one.
+
+        Args:
+            other: The NodeDependencies object to aggregate from
+        """
+        # Aggregate all dependency types - no None checks needed since we use default_factory=set
+        self.referenced_workflows.update(other.referenced_workflows)
+        self.static_files.update(other.static_files)
+        self.imports.update(other.imports)
+        self.libraries.update(other.libraries)
+
 
 class NodeResolutionState(StrEnum):
     """Possible states for a node during resolution."""
@@ -59,6 +108,23 @@ class NodeResolutionState(StrEnum):
     RESOLVED = auto()
 
 
+def get_library_names_with_publish_handlers() -> list[str]:
+    """Get names of all registered libraries that have PublishWorkflowRequest handlers."""
+    from griptape_nodes.retained_mode.events.workflow_events import PublishWorkflowRequest
+    from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+    library_manager = GriptapeNodes.LibraryManager()
+    event_handlers = library_manager.get_registered_event_handlers(PublishWorkflowRequest)
+
+    # Always include "local" as the first option
+    library_names = [LOCAL_EXECUTION]
+
+    # Add all registered library names that can handle PublishWorkflowRequest
+    library_names.extend(sorted(event_handlers.keys()))
+
+    return library_names
+
+
 class BaseNode(ABC):
     # Owned by a flow
     name: str
@@ -104,6 +170,16 @@ class BaseNode(ABC):
         self.process_generator = None
         self._tracked_parameters = []
         self.set_entry_control_parameter(None)
+        self.execution_environment = Parameter(
+            name="execution_environment",
+            tooltip="Environment that the node should execute in",
+            type=ParameterTypeBuiltin.STR,
+            allowed_modes={ParameterMode.PROPERTY},
+            default_value=LOCAL_EXECUTION,
+            traits={Options(choices=get_library_names_with_publish_handlers())},
+            ui_options={"hide": True},
+        )
+        self.add_parameter(self.execution_environment)
 
     # This is gross and we need to have a universal pass on resolution state changes and emission of events. That's what this ticket does!
     # https://github.com/griptape-ai/griptape-nodes/issues/994
@@ -384,9 +460,7 @@ class BaseNode(ABC):
         for name in names:
             parameter = self.get_parameter_by_name(name)
             if parameter is not None:
-                ui_options = parameter.ui_options
-                ui_options["hide"] = not visible
-                parameter.ui_options = ui_options
+                parameter.ui_options = {**parameter.ui_options, "hide": not visible}
 
     def get_message_by_name_or_element_id(self, element: str) -> ParameterMessage | None:
         element_items = self.root_ui_element.find_elements_by_type(ParameterMessage)
@@ -408,9 +482,7 @@ class BaseNode(ABC):
         for name in names:
             message = self.get_message_by_name_or_element_id(name)
             if message is not None:
-                ui_options = message.ui_options
-                ui_options["hide"] = not visible
-                message.ui_options = ui_options
+                message.ui_options = {**message.ui_options, "hide": not visible}
 
     def hide_message_by_name(self, names: str | list[str]) -> None:
         self._set_message_visibility(names, visible=False)
@@ -727,11 +799,9 @@ class BaseNode(ABC):
                 return param
         return None
 
-    # Abstract method to process the node. Must be defined by the type
     # Must save the values of the output parameters in NodeContext.
-    @abstractmethod
-    def process[T](self) -> AsyncResult | None:
-        pass
+    def process(self) -> AsyncResult | None:
+        raise NotImplementedError
 
     async def aprocess(self) -> None:
         """Async version of process().
@@ -825,6 +895,35 @@ class BaseNode(ABC):
         # Then clear the reference to the first spotlight parameter
         self.current_spotlight_parameter = None
 
+    def get_node_dependencies(self) -> NodeDependencies | None:
+        """Return the dependencies that this node has on external resources.
+
+        This method should be overridden by nodes that have dependencies on:
+        - Referenced workflows: Other workflows that this node calls or references
+        - Static files: Files that this node reads from or requires for operation
+        - Python imports: Modules or classes that this node imports beyond standard dependencies
+
+        This information can be used by the system for workflow packaging, dependency
+        resolution, deployment planning, and ensuring all required resources are available.
+
+        Returns:
+            NodeDependencies object containing the node's dependencies, or None if the node
+            has no external dependencies beyond the standard framework dependencies.
+
+        Example:
+            def get_node_dependencies(self) -> NodeDependencies | None:
+                return NodeDependencies(
+                    referenced_workflows={"image_processing_workflow", "validation_workflow"},
+                    static_files={"config.json", "model_weights.pkl"},
+                    imports={
+                        ImportDependency("numpy"),
+                        ImportDependency("sklearn.linear_model", "LinearRegression"),
+                        ImportDependency("custom_module", "SpecialProcessor")
+                    }
+                )
+        """
+        return None
+
     def append_value_to_parameter(self, parameter_name: str, value: Any) -> None:
         from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
 
@@ -867,7 +966,7 @@ class BaseNode(ABC):
             msg = f"Parameter '{parameter_name} doesn't exist on {self.name}'"
             raise RuntimeError(msg)
 
-    def reorder_elements(self, element_order: list[str | int]) -> None:
+    def reorder_elements(self, element_order: list[str] | list[int] | list[str | int]) -> None:
         """Reorder the elements of this node.
 
         Args:
@@ -1109,7 +1208,10 @@ class TrackedParameterOutputValues(dict[str, Any]):
             keys_to_clear = list(self.keys())
             super().clear()
             for key in keys_to_clear:
-                self._emit_parameter_change_event(key, None, deleted=True)
+                # Some nodes still have values set, even if their output values are cleared
+                # Here, we are emitting an event with those set values, to not misrepresent the values of the parameters in the UI.
+                value = self._node.get_parameter_value(key)
+                self._emit_parameter_change_event(key, value, deleted=True)
 
     def silent_clear(self) -> None:
         """Clear all values without emitting parameter change events."""
@@ -1393,6 +1495,16 @@ class EndNode(BaseNode):
 
         self.status_component.set_execution_result(was_successful=was_successful, result_details=details)
 
+        # Update all values to use the output value
+        for param in self.parameters:
+            if param.type != ParameterTypeBuiltin.CONTROL_TYPE:
+                value = self.get_parameter_value(param.name)
+                self.parameter_output_values[param.name] = value
+        next_control_output = self.get_next_control_output()
+        # Update which control parameter to flag as the output value.
+        if next_control_output is not None:
+            self.parameter_output_values[next_control_output.name] = 1
+
 
 class StartLoopNode(BaseNode):
     end_node: EndLoopNode | None = None

griptape_nodes/machines/control_flow.py

@@ -15,6 +15,7 @@ from griptape_nodes.retained_mode.events.base_events import ExecutionEvent, Exec
 from griptape_nodes.retained_mode.events.execution_events import (
     ControlFlowResolvedEvent,
     CurrentControlNodeEvent,
+    InvolvedNodesEvent,
     SelectedControlOutputEvent,
 )
 from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
@@ -252,12 +253,21 @@ class ControlFlowMachine(FSM[ControlFlowContext]):
             current_nodes = await self._process_nodes_for_dag(start_node)
         else:
             current_nodes = [start_node]
+            # For control flow/sequential: emit all nodes in flow as involved
         self._context.current_nodes = current_nodes
         # Set entry control parameter for initial node (None for workflow start)
         for node in current_nodes:
             node.set_entry_control_parameter(None)
         # Set up to debug
         self._context.paused = debug_mode
+        flow_manager = GriptapeNodes.FlowManager()
+        flow = flow_manager.get_flow_by_name(self._context.flow_name)
+        involved_nodes = list(flow.nodes.keys())
+        GriptapeNodes.EventManager().put_event(
+            ExecutionGriptapeNodeEvent(
+                wrapped_event=ExecutionEvent(payload=InvolvedNodesEvent(involved_nodes=involved_nodes))
+            )
+        )
         await self.start(ResolveNodeState)  # Begins the flow
 
     async def update(self) -> None:

griptape_nodes/machines/dag_builder.py

@@ -43,10 +43,12 @@ class DagBuilder:
 
     graphs: dict[str, DirectedGraph]  # Str is the name of the start node associated here.
     node_to_reference: dict[str, DagNode]
+    graph_to_nodes: dict[str, set[str]]  # Track which nodes belong to which graph
 
     def __init__(self) -> None:
         self.graphs = {}
         self.node_to_reference: dict[str, DagNode] = {}
+        self.graph_to_nodes = {}
 
     # 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
@@ -59,16 +61,15 @@ class DagBuilder:
         if graph is None:
             graph = DirectedGraph()
             self.graphs[graph_name] = graph
+            self.graph_to_nodes[graph_name] = set()
 
         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.
@@ -98,6 +99,10 @@ class DagBuilder:
             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)
+
+            # Track which nodes belong to this graph
+            self.graph_to_nodes[graph_name].add(current_node.name)
+
             # DON'T mark as resolved - that happens during actual execution
             added_nodes.append(current_node)
 
@@ -117,12 +122,19 @@ class DagBuilder:
             graph = DirectedGraph()
             self.graphs[graph_name] = graph
         graph.add_node(node_for_adding=node.name)
+
+        # Track which nodes belong to this graph
+        if graph_name not in self.graph_to_nodes:
+            self.graph_to_nodes[graph_name] = set()
+        self.graph_to_nodes[graph_name].add(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()
+        self.graph_to_nodes.clear()
 
     def can_queue_control_node(self, node: DagNode) -> bool:
         if len(self.graphs) == 1:
@@ -205,3 +217,10 @@ class DagBuilder:
                                 return True
 
         return False
+
+    def cleanup_empty_graph_nodes(self, graph_name: str) -> None:
+        """Remove nodes from node_to_reference when their graph becomes empty (only in single node resolution)."""
+        if graph_name in self.graph_to_nodes:
+            for node_name in self.graph_to_nodes[graph_name]:
+                self.node_to_reference.pop(node_name, None)
+            self.graph_to_nodes.pop(graph_name, None)