diff --git a/OpenRA.Mods.Common/Traits/Cloak.cs b/OpenRA.Mods.Common/Traits/Cloak.cs
index b875ef056a..3e5706a8b8 100644
--- a/OpenRA.Mods.Common/Traits/Cloak.cs
+++ b/OpenRA.Mods.Common/Traits/Cloak.cs
@@ -54,8 +54,8 @@ namespace OpenRA.Mods.Common.Traits
 		public readonly HashSet<string> CloakTypes = new HashSet<string> { "Cloak" };
 
 		[UpgradeGrantedReference]
-		[Desc("The upgrades to grant to self while cloaked.")]
-		public readonly string[] WhileCloakedUpgrades = { };
+		[Desc("The condition to grant to self while cloaked.")]
+		public readonly string CloakedCondition = null;
 
 		public override object Create(ActorInitializer init) { return new Cloak(this); }
 	}
@@ -70,6 +70,7 @@ namespace OpenRA.Mods.Common.Traits
 
 		CPos? lastPos;
 		bool wasCloaked = false;
+		int cloakedToken = UpgradeManager.InvalidConditionToken;
 
 		public Cloak(CloakInfo info)
 			: base(info)
@@ -81,12 +82,11 @@ namespace OpenRA.Mods.Common.Traits
 		{
 			upgradeManager = self.TraitOrDefault<UpgradeManager>();
 
-			// The upgrade manager exists, but may not have finished being created yet.
-			// We'll defer the upgrades until the end of the tick, at which point it will be ready.
 			if (Cloaked)
 			{
 				wasCloaked = true;
-				self.World.AddFrameEndTask(_ => GrantUpgrades(self));
+				if (upgradeManager != null && cloakedToken == UpgradeManager.InvalidConditionToken && !string.IsNullOrEmpty(Info.CloakedCondition))
+					cloakedToken = upgradeManager.GrantCondition(self, Info.CloakedCondition);
 			}
 		}
 
@@ -144,13 +144,17 @@ namespace OpenRA.Mods.Common.Traits
 			var isCloaked = Cloaked;
 			if (isCloaked && !wasCloaked)
 			{
-				GrantUpgrades(self);
+				if (upgradeManager != null && cloakedToken == UpgradeManager.InvalidConditionToken && !string.IsNullOrEmpty(Info.CloakedCondition))
+					cloakedToken = upgradeManager.GrantCondition(self, Info.CloakedCondition);
+
 				if (!self.TraitsImplementing<Cloak>().Any(a => a != this && a.Cloaked))
 					Game.Sound.Play(Info.CloakSound, self.CenterPosition);
 			}
 			else if (!isCloaked && wasCloaked)
 			{
-				RevokeUpgrades(self);
+				if (cloakedToken != UpgradeManager.InvalidConditionToken)
+					cloakedToken = upgradeManager.RevokeCondition(self, cloakedToken);
+
 				if (!self.TraitsImplementing<Cloak>().Any(a => a != this && a.Cloaked))
 					Game.Sound.Play(Info.UncloakSound, self.CenterPosition);
 			}
@@ -176,20 +180,6 @@ namespace OpenRA.Mods.Common.Traits
 			return color;
 		}
 
-		void GrantUpgrades(Actor self)
-		{
-			if (upgradeManager != null)
-				foreach (var u in Info.WhileCloakedUpgrades)
-					upgradeManager.GrantUpgrade(self, u, this);
-		}
-
-		void RevokeUpgrades(Actor self)
-		{
-			if (upgradeManager != null)
-				foreach (var u in Info.WhileCloakedUpgrades)
-					upgradeManager.RevokeUpgrade(self, u, this);
-		}
-
 		void INotifyHarvesterAction.MovingToResources(Actor self, CPos targetCell, Activity next) { }
 
 		void INotifyHarvesterAction.MovingToRefinery(Actor self, CPos targetCell, Activity next) { }
diff --git a/OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs b/OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs
index bd8edfa240..373232e9f1 100644
--- a/OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs
+++ b/OpenRA.Mods.Common/UtilityCommands/UpgradeRules.cs
@@ -593,6 +593,9 @@ namespace OpenRA.Mods.Common.UtilityCommands
 						ConvertUpgradesToCondition(parent, node, "AirborneUpgrades", "AirborneCondition");
 						ConvertUpgradesToCondition(parent, node, "CruisingUpgrades", "CruisingCondition");
 					}
+
+					if (node.Key.StartsWith("Cloak", StringComparison.Ordinal))
+						ConvertUpgradesToCondition(parent, node, "WhileCloakedUpgrades", "CloakedCondition");
 				}
 
 				UpgradeActorRules(modData, engineVersion, ref node.Value.Nodes, node, depth + 1);
diff --git a/mods/ra/rules/ships.yaml b/mods/ra/rules/ships.yaml
index 7af2caa9d0..d5b48e149f 100644
--- a/mods/ra/rules/ships.yaml
+++ b/mods/ra/rules/ships.yaml
@@ -31,7 +31,7 @@ SS:
 		CloakDelay: 50
 		CloakSound: subshow1.aud
 		UncloakSound: subshow1.aud
-		WhileCloakedUpgrades: underwater
+		CloakedCondition: underwater
 		Palette: submerged
 	Armament:
 		Weapon: TorpTube
@@ -85,7 +85,7 @@ MSUB:
 		CloakDelay: 100
 		CloakSound: subshow1.aud
 		UncloakSound: subshow1.aud
-		WhileCloakedUpgrades: underwater
+		CloakedCondition: underwater
 		Palette: submerged
 	Armament@PRIMARY:
 		Weapon: SubMissile