diff --git a/OpenRA.Mods.Common/Traits/BotModules/Squads/States/StateBase.cs b/OpenRA.Mods.Common/Traits/BotModules/Squads/States/StateBase.cs
index 983f65a052..3c5c23f20a 100644
--- a/OpenRA.Mods.Common/Traits/BotModules/Squads/States/StateBase.cs
+++ b/OpenRA.Mods.Common/Traits/BotModules/Squads/States/StateBase.cs
@@ -106,21 +106,7 @@ namespace OpenRA.Mods.Common.Traits.BotModules.Squads
 
 		protected static bool IsRearming(Actor a)
 		{
-			if (a.IsIdle)
-				return false;
-
-			var activity = a.CurrentActivity;
-			if (activity.GetType() == typeof(Resupply))
-				return true;
-
-			var next = activity.NextActivity;
-			if (next == null)
-				return false;
-
-			if (next.GetType() == typeof(Resupply))
-				return true;
-
-			return false;
+			return !a.IsIdle && (a.CurrentActivity.ActivitiesImplementing<Resupply>().Any() || a.CurrentActivity.ActivitiesImplementing<ReturnToBase>().Any());
 		}
 
 		protected static bool FullAmmo(IEnumerable<AmmoPool> ammoPools)