PLT-2761 - Prevent ID change for initial nodes and fix edge ids (#1999)

Author: paulnoirel

libs/labelbox/src/labelbox/schema/workflow/edges.py

@@ -5,7 +5,6 @@
 """
 
 import logging
-import uuid
 from typing import Dict, Any, Optional, TYPE_CHECKING
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
 
@@ -25,11 +24,22 @@ class WorkflowEdge(BaseModel):
     from a source node to a target node through specific handles.
 
     Attributes:
-        id: Unique identifier for the edge
+        id: Unique identifier for the edge (format: xy-edge__{source}{sourceHandle}-{target}{targetHandle})
         source: ID of the source node
         target: ID of the target node
-        sourceHandle: Output handle on the source node (e.g., 'if', 'else')
+        sourceHandle: Output handle on the source node (e.g., 'if', 'else', 'approved', 'rejected')
         targetHandle: Input handle on the target node (typically 'in')
+
+    Edge ID Format:
+        Edge IDs follow the pattern: xy-edge__{source}{sourceHandle}-{target}{targetHandle}
+
+        Example: xy-edge__node1if-node2in
+        - Prefix: xy-edge__
+        - Source node ID: node1
+        - Source handle: if
+        - Separator: -
+        - Target node ID: node2
+        - Target handle: in
     """
 
     id: str
@@ -220,7 +230,13 @@ def _create_edge_instance(
         Returns:
             Created WorkflowEdge instance
         """
-        edge_id = f"edge-{uuid.uuid4()}"
+        # Generate edge ID using the correct format: xy-edge__{source}{sourceHandle}-{target}{targetHandle}
+        source_handle = output_type.value
+        target_handle = "in"
+        edge_id = (
+            f"xy-edge__{source.id}{source_handle}-{target.id}{target_handle}"
+        )
+
         logger.debug(
             f"Creating edge {edge_id} from {source.id} to {target.id} with type {output_type.value}"
         )
@@ -229,8 +245,8 @@ def _create_edge_instance(
             id=edge_id,
             source=source.id,
             target=target.id,
-            sourceHandle=output_type.value,
-            targetHandle="in",  # Explicitly set targetHandle
+            sourceHandle=source_handle,
+            targetHandle=target_handle,  # Explicitly set targetHandle
         )
         edge.set_workflow_reference(self.workflow)
         return edge

libs/labelbox/src/labelbox/schema/workflow/workflow.py

@@ -496,7 +496,10 @@ def reset_to_initial_nodes(
     ) -> InitialNodes:
         """Reset workflow and create the two required initial nodes.
 
-        Clears all existing nodes and edges, then creates:
+        IMPORTANT: This method preserves existing initial node IDs to prevent workflow breakage.
+        It only creates new IDs for truly new workflows (first-time setup).
+
+        Clears all non-initial nodes and edges, then creates/updates:
         - InitialLabeling node: Entry point for new data requiring labeling
         - InitialRework node: Entry point for rejected data requiring corrections
 
@@ -526,11 +529,28 @@ def reset_to_initial_nodes(
             rework_config.model_dump(exclude_none=True) if rework_config else {}
         )
 
-        # Reset workflow configuration
+        # Find existing initial nodes to preserve their IDs
+        existing_labeling_id = None
+        existing_rework_id = None
+
+        for node_data in self.config.get("nodes", []):
+            definition_id = node_data.get("definitionId")
+            if definition_id == WorkflowDefinitionId.InitialLabelingTask.value:
+                existing_labeling_id = node_data.get("id")
+            elif definition_id == WorkflowDefinitionId.InitialReworkTask.value:
+                existing_rework_id = node_data.get("id")
+
+        # Reset workflow configuration (clear all nodes and edges)
         self.config = {"nodes": [], "edges": []}
         self._nodes_cache = None
         self._edges_cache = None
 
+        # Create/recreate initial nodes, preserving existing IDs if they exist
+        if existing_labeling_id:
+            labeling_dict["id"] = existing_labeling_id
+        if existing_rework_id:
+            rework_dict["id"] = existing_rework_id
+
         # Create required initial nodes using internal method
         initial_labeling = cast(
             InitialLabelingNode,
@@ -554,15 +574,23 @@ def copy_workflow_structure(
         target_client,
         target_project_id: str,
     ) -> "ProjectWorkflow":
-        """Copy the workflow structure from a source workflow to a new project."""
+        """Copy the workflow structure from a source workflow to a new project.
+
+        IMPORTANT: This method preserves existing initial node IDs to prevent workflow breakage.
+        Changing initial node IDs will completely break the workflow and require support intervention.
+        """
         return WorkflowOperations.copy_workflow_structure(
             source_workflow, target_client, target_project_id
         )
 
     def copy_from(
         self, source_workflow: "ProjectWorkflow", auto_layout: bool = True
     ) -> "ProjectWorkflow":
-        """Copy the nodes and edges from a source workflow to this workflow."""
+        """Copy the nodes and edges from a source workflow to this workflow.
+
+        IMPORTANT: This method preserves existing initial node IDs to prevent workflow breakage.
+        Changing initial node IDs will completely break the workflow and require support intervention.
+        """
         return WorkflowOperations.copy_from(self, source_workflow, auto_layout)
 
     # Layout and display methods

libs/labelbox/src/labelbox/schema/workflow/workflow_operations.py

@@ -390,7 +390,11 @@ def copy_workflow_structure(
         target_client,
         target_project_id: str,
     ) -> "ProjectWorkflow":
-        """Copy the workflow structure from a source workflow to a new project."""
+        """Copy the workflow structure from a source workflow to a new project.
+
+        IMPORTANT: This method preserves existing initial node IDs in the target workflow
+        to prevent workflow breakage. Only non-initial nodes get new IDs.
+        """
         try:
             # Create a new workflow in the target project
             from labelbox.schema.workflow.workflow import ProjectWorkflow
@@ -399,38 +403,74 @@ def copy_workflow_structure(
                 target_client, target_project_id
             )
 
+            # Find existing initial nodes in target workflow to preserve their IDs
+            existing_initial_ids = {}
+            for node_data in target_workflow.config.get("nodes", []):
+                definition_id = node_data.get("definitionId")
+                if (
+                    definition_id
+                    == WorkflowDefinitionId.InitialLabelingTask.value
+                ):
+                    existing_initial_ids[
+                        WorkflowDefinitionId.InitialLabelingTask.value
+                    ] = node_data.get("id")
+                elif (
+                    definition_id
+                    == WorkflowDefinitionId.InitialReworkTask.value
+                ):
+                    existing_initial_ids[
+                        WorkflowDefinitionId.InitialReworkTask.value
+                    ] = node_data.get("id")
+
             # Get the source config
             new_config = source_workflow.config.copy()
             old_to_new_id_map = {}
 
-            # Generate new IDs for all nodes
+            # Generate new IDs for all nodes, but preserve existing initial node IDs
             if new_config.get("nodes"):
-                new_config["nodes"] = [
-                    {
-                        **node,
-                        "id": str(uuid.uuid4()),
-                    }
-                    for node in new_config["nodes"]
-                ]
-                # Create mapping of old to new IDs
-                old_to_new_id_map = {
-                    old_node["id"]: new_node["id"]
-                    for old_node, new_node in zip(
-                        source_workflow.config["nodes"], new_config["nodes"]
+                updated_nodes = []
+                for node in new_config["nodes"]:
+                    definition_id = node.get("definitionId")
+                    old_id = node["id"]
+
+                    # Preserve existing initial node IDs, generate new IDs for others
+                    if definition_id in existing_initial_ids:
+                        new_id = existing_initial_ids[definition_id]
+                    else:
+                        new_id = str(uuid.uuid4())
+
+                    old_to_new_id_map[old_id] = new_id
+                    updated_nodes.append(
+                        {
+                            **node,
+                            "id": new_id,
+                        }
                     )
-                }
+
+                new_config["nodes"] = updated_nodes
 
             # Update edges to use the new node IDs
             if new_config.get("edges"):
-                new_config["edges"] = [
-                    {
-                        **edge,
-                        "id": str(uuid.uuid4()),
-                        "source": old_to_new_id_map[edge["source"]],
-                        "target": old_to_new_id_map[edge["target"]],
-                    }
-                    for edge in new_config["edges"]
-                ]
+                updated_edges = []
+                for edge in new_config["edges"]:
+                    source_id = old_to_new_id_map[edge["source"]]
+                    target_id = old_to_new_id_map[edge["target"]]
+                    source_handle = edge.get("sourceHandle", "if")
+                    target_handle = edge.get("targetHandle", "in")
+
+                    # Generate edge ID using correct format: xy-edge__{source}{sourceHandle}-{target}{targetHandle}
+                    edge_id = f"xy-edge__{source_id}{source_handle}-{target_id}{target_handle}"
+
+                    updated_edges.append(
+                        {
+                            **edge,
+                            "id": edge_id,
+                            "source": source_id,
+                            "target": target_id,
+                        }
+                    )
+
+                new_config["edges"] = updated_edges
 
             # Update the target workflow with the new config
             target_workflow.config = new_config
@@ -450,8 +490,31 @@ def copy_from(
         source_workflow: "ProjectWorkflow",
         auto_layout: bool = True,
     ) -> "ProjectWorkflow":
-        """Copy the nodes and edges from a source workflow to this workflow."""
+        """Copy the nodes and edges from a source workflow to this workflow.
+
+        IMPORTANT: This method preserves existing initial node IDs in the target workflow
+        to prevent workflow breakage. Only non-initial nodes get new IDs.
+        """
         try:
+            # Find existing initial nodes in target workflow to preserve their IDs
+            existing_initial_ids = {}
+            for node_data in workflow.config.get("nodes", []):
+                definition_id = node_data.get("definitionId")
+                if (
+                    definition_id
+                    == WorkflowDefinitionId.InitialLabelingTask.value
+                ):
+                    existing_initial_ids[
+                        WorkflowDefinitionId.InitialLabelingTask.value
+                    ] = node_data.get("id")
+                elif (
+                    definition_id
+                    == WorkflowDefinitionId.InitialReworkTask.value
+                ):
+                    existing_initial_ids[
+                        WorkflowDefinitionId.InitialReworkTask.value
+                    ] = node_data.get("id")
+
             # Create a clean work config (without connections)
             work_config: Dict[str, List[Any]] = {"nodes": [], "edges": []}
 
@@ -463,9 +526,15 @@ def copy_from(
 
             # First pass: Create all nodes by directly copying configuration
             for source_node_data in source_workflow.config.get("nodes", []):
-                # Generate a new ID for the node
-                new_id = f"node-{uuid.uuid4()}"
+                definition_id = source_node_data.get("definitionId")
                 old_id = source_node_data.get("id")
+
+                # Preserve existing initial node IDs, generate new IDs for others
+                if definition_id in existing_initial_ids:
+                    new_id = existing_initial_ids[definition_id]
+                else:
+                    new_id = f"node-{uuid.uuid4()}"
+
                 id_mapping[old_id] = new_id
 
                 # Create a new node data dictionary by copying the source node
@@ -498,12 +567,18 @@ def copy_from(
                     continue
 
                 # Create new edge
+                source_handle = source_edge_data.get("sourceHandle", "out")
+                target_handle = source_edge_data.get("targetHandle", "in")
+
+                # Generate edge ID using correct format: xy-edge__{source}{sourceHandle}-{target}{targetHandle}
+                edge_id = f"xy-edge__{id_mapping[source_id]}{source_handle}-{id_mapping[target_id]}{target_handle}"
+
                 new_edge = {
-                    "id": f"edge-{uuid.uuid4()}",
+                    "id": edge_id,
                     "source": id_mapping[source_id],
                     "target": id_mapping[target_id],
-                    "sourceHandle": source_edge_data.get("sourceHandle", "out"),
-                    "targetHandle": source_edge_data.get("targetHandle", "in"),
+                    "sourceHandle": source_handle,
+                    "targetHandle": target_handle,
                 }
 
                 # Add the edge to config