diff --git a/OpenRA.Game/Scripting/ScriptContext.cs b/OpenRA.Game/Scripting/ScriptContext.cs
index 499de38637..f869fb2820 100644
--- a/OpenRA.Game/Scripting/ScriptContext.cs
+++ b/OpenRA.Game/Scripting/ScriptContext.cs
@@ -254,9 +254,9 @@ namespace OpenRA.Scripting
 		static readonly object[] NoArguments = new object[0];
 		Type[] FilterActorCommands(ActorInfo ai)
 		{
-			var method = typeof(TypeDictionary).GetMethod("Contains");
+			var method = typeof(ActorInfo).GetMethod("HasTraitInfo");
 			return knownActorCommands.Where(c => ExtractRequiredTypes(c)
-				.All(t => (bool)method.MakeGenericMethod(t).Invoke(ai.Traits, NoArguments)))
+				.All(t => (bool)method.MakeGenericMethod(t).Invoke(ai, NoArguments)))
 				.ToArray();
 		}
 
diff --git a/OpenRA.Mods.Common/AI/BaseBuilder.cs b/OpenRA.Mods.Common/AI/BaseBuilder.cs
index 46a102c75d..5f57064ff2 100644
--- a/OpenRA.Mods.Common/AI/BaseBuilder.cs
+++ b/OpenRA.Mods.Common/AI/BaseBuilder.cs
@@ -141,9 +141,9 @@ namespace OpenRA.Mods.Common.AI
 				// HACK: HACK HACK HACK
 				// TODO: Derive this from BuildingCommonNames instead
 				var type = BuildingType.Building;
-				if (world.Map.Rules.Actors[currentBuilding.Item].Traits.Contains<AttackBaseInfo>())
+				if (world.Map.Rules.Actors[currentBuilding.Item].HasTraitInfo<AttackBaseInfo>())
 					type = BuildingType.Defense;
-				else if (world.Map.Rules.Actors[currentBuilding.Item].Traits.Contains<RefineryInfo>())
+				else if (world.Map.Rules.Actors[currentBuilding.Item].HasTraitInfo<RefineryInfo>())
 					type = BuildingType.Refinery;
 
 				var location = ai.ChooseBuildLocation(currentBuilding.Item, true, type);