Stage 4: Clean entity merging on upload with DraftContext support

Add draft_entity_info parameter to set_up_params_for_uploading() to use
DraftContext's entity_info as source of truth during upload. Implement
_merge_draft_entities_from_params() to collect draft entities from params
and merge into entity_info without duplicates.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <[email protected]>

Author: benflexcompute

flow360/component/project.py

@@ -32,7 +32,10 @@
     validate_params_with_context,
 )
 from flow360.component.resource_base import Flow360Resource
-from flow360.component.simulation.draft_context.context import DraftContext
+from flow360.component.simulation.draft_context.context import (
+    DraftContext,
+    get_active_draft,
+)
 from flow360.component.simulation.entity_info import GeometryEntityInfo
 from flow360.component.simulation.folder import Folder
 from flow360.component.simulation.simulation_params import SimulationParams
@@ -1480,12 +1483,17 @@ def _run(
         if use_geometry_AI is True and use_beta_mesher is False:
             raise Flow360ValueError("Enabling Geometry AI requires also enabling beta mesher.")
 
+        # Check if there's an active DraftContext and get its entity_info
+        active_draft = get_active_draft()
+        draft_entity_info = active_draft._entity_info if active_draft is not None else None
+
         params = set_up_params_for_uploading(
             params=params,
             root_asset=self._root_asset,
             length_unit=self.length_unit,
             use_beta_mesher=use_beta_mesher,
             use_geometry_AI=use_geometry_AI,
+            draft_entity_info=draft_entity_info,
         )
 
         params, errors = validate_params_with_context(

flow360/component/project_utils.py

@@ -255,6 +255,9 @@ def _set_up_params_non_persistent_entity_info(entity_info, params: SimulationPar
     """
     Setting up non-persistent entities (AKA draft entities) in params.
     Add the ones used to the entity info.
+
+    LEGACY: This function is used for the legacy workflow (without DraftContext).
+    For DraftContext workflow, use _merge_draft_entities_from_params() instead.
     """
 
     entity_registry = params.used_entity_registry
@@ -270,6 +273,49 @@ def _set_up_params_non_persistent_entity_info(entity_info, params: SimulationPar
     return entity_info
 
 
+def _merge_draft_entities_from_params(
+    entity_info: EntityInfoModel,
+    params: SimulationParams,
+) -> EntityInfoModel:
+    """
+    Collect draft entities from params.used_entity_registry and merge into entity_info.
+
+    This function implements the merging logic for the DraftContext workflow:
+    - If a draft entity already exists in entity_info (by ID), use entity_info version (source of truth)
+    - If a draft entity is new (not in entity_info), add it from params
+
+    This ensures that:
+    1. Entities managed by DraftContext retain their modifications
+    2. New entities created by the user during simulation setup are captured
+
+    Parameters:
+        entity_info: The entity_info to merge into (typically from DraftContext)
+        params: The SimulationParams containing used_entity_registry
+
+    Returns:
+        EntityInfoModel: The updated entity_info with merged draft entities
+    """
+    used_registry = params.used_entity_registry
+
+    # Get all draft entity types from the DraftEntityTypes annotation
+    draft_type_union = get_args(DraftEntityTypes)[0]
+    draft_type_list = get_args(draft_type_union)
+
+    # Build a set of IDs already in entity_info for quick lookup (Draft entities have unique UUIDs)
+    existing_ids = {e.private_attribute_id for e in entity_info.draft_entities}
+
+    for draft_type in draft_type_list:
+        draft_entities_used = list(used_registry.view(draft_type))
+        for draft_entity in draft_entities_used:
+            # Only add if not already in entity_info (by ID)
+            # If already present, entity_info version is source of truth - keep it as is
+            if draft_entity.private_attribute_id not in existing_ids:
+                entity_info.draft_entities.append(draft_entity)
+                existing_ids.add(draft_entity.private_attribute_id)
+
+    return entity_info
+
+
 def _update_entity_grouping_tags(entity_info, params: SimulationParams) -> EntityInfoModel:
     """
     Update the entity grouping tags in params to resolve possible conflicts
@@ -411,9 +457,20 @@ def set_up_params_for_uploading(
     params: SimulationParams,
     use_beta_mesher: bool,
     use_geometry_AI: bool,  # pylint: disable=invalid-name
+    draft_entity_info: Optional[EntityInfoModel] = None,
 ) -> SimulationParams:
     """
     Set up params before submitting the draft.
+
+    Parameters:
+        root_asset: The root asset (Geometry, SurfaceMesh, or VolumeMesh).
+        length_unit: The project length unit.
+        params: The SimulationParams to set up.
+        use_beta_mesher: Whether to use the beta mesher.
+        use_geometry_AI: Whether to use Geometry AI.
+        draft_entity_info: Optional entity_info from DraftContext. When provided,
+            this is used as the source of truth for entities instead of root_asset.entity_info (legacy behavior).
+            This enables proper entity isolation for the DraftContext workflow.
     """
 
     with model_attribute_unlock(params.private_attribute_asset_cache, "project_length_unit"):
@@ -429,18 +486,28 @@ def set_up_params_for_uploading(
             use_geometry_AI if use_geometry_AI else False
         )
 
-    # User may have made modifications to the entities which is recorded in asset's entity registry
-    # We need to reflect these changes.
-    root_asset.entity_info.update_persistent_entities(
-        asset_entity_registry=root_asset.internal_registry
-    )
+    if draft_entity_info is not None:
+        # New DraftContext workflow: use draft's entity_info as source of truth
+        # Merge draft entities from params.used_entity_registry into draft_entity_info
+        entity_info = _merge_draft_entities_from_params(draft_entity_info, params)
+
+        # Update entity grouping tags if needed
+        # (back compatibility, since the grouping should already have been captured in the draft_entity_info)
+        entity_info = _update_entity_grouping_tags(entity_info, params)
+    else:
+        # Legacy workflow (without DraftContext): use root_asset.entity_info
+        # User may have made modifications to the entities which is recorded in asset's entity registry
+        # We need to reflect these changes.
+        root_asset.entity_info.update_persistent_entities(
+            asset_entity_registry=root_asset.internal_registry
+        )
 
-    # Check if there are any new draft entities that have been added in the params by the user
-    entity_info = _set_up_params_non_persistent_entity_info(root_asset.entity_info, params)
+        # Check if there are any new draft entities that have been added in the params by the user
+        entity_info = _set_up_params_non_persistent_entity_info(root_asset.entity_info, params)
 
-    # If the customer just load the param without re-specify the same set of entity grouping tags,
-    # we need to update the entity grouping tags to the ones in the SimulationParams.
-    entity_info = _update_entity_grouping_tags(entity_info, params)
+        # If the customer just load the param without re-specify the same set of entity grouping tags,
+        # we need to update the entity grouping tags to the ones in the SimulationParams.
+        entity_info = _update_entity_grouping_tags(entity_info, params)
 
     with model_attribute_unlock(params.private_attribute_asset_cache, "project_entity_info"):
         params.private_attribute_asset_cache.project_entity_info = entity_info

flow360/component/simulation/draft_context/context.py

@@ -102,6 +102,7 @@ def __exit__(self, exc_type, exc, exc_tb) -> None:
     # endregion ------------------------------------------------------------------------------------
 
     # region -----------------------------Public properties Below-------------------------------------
+
     # Persistent entities
     @property
     def body_groups(self) -> EntityRegistryView: