diff --git a/Gems/Multiplayer/Code/Source/NetworkEntity/NetworkEntityManager.cpp b/Gems/Multiplayer/Code/Source/NetworkEntity/NetworkEntityManager.cpp index db7f7243cc..b72ea21719 100644 --- a/Gems/Multiplayer/Code/Source/NetworkEntity/NetworkEntityManager.cpp +++ b/Gems/Multiplayer/Code/Source/NetworkEntity/NetworkEntityManager.cpp @@ -21,6 +21,8 @@ #include #include #include +#include +#include #include namespace Multiplayer @@ -272,15 +274,32 @@ namespace Multiplayer bool safeToExit = true; NetworkEntityHandle entityHandle = m_networkEntityTracker.Get(entityId); - // We also need special handling for the EntityHierarchyComponent as well, since related entities need to be migrated together - //auto* hierarchyController = FindController(nonConstExitingEntityPtr); - //if (hierarchyController) - //{ - // if (hierarchyController->GetParentRelatedEntity()) - // { - // safeToExit = false; - // } - //} + // We also need special handling for the NetworkHierarchy as well, since related entities need to be migrated together + NetworkHierarchyRootComponentController* hierarchyRootController = entityHandle.FindController(); + NetworkHierarchyChildComponentController* hierarchyChildController = entityHandle.FindController(); + + // Find the root entity + AZ::Entity* hierarchyRootEntity = nullptr; + if (hierarchyRootController) + { + hierarchyRootEntity = hierarchyRootController->GetParent().GetHierarchicalRoot(); + } + else if (hierarchyChildController) + { + hierarchyRootEntity = hierarchyChildController->GetParent().GetHierarchicalRoot(); + } + + if (hierarchyRootEntity) + { + NetEntityId rootNetId = GetNetEntityIdById(hierarchyRootEntity->GetId()); + ConstNetworkEntityHandle rootEntityHandle = GetEntity(rootNetId); + + // Check if the root entity is still tracked by this authority + if (rootEntityHandle.Exists() && rootEntityHandle.GetNetBindComponent()->HasController()) + { + safeToExit = false; + } + } // Validate that we aren't already planning to remove this entity if (safeToExit)