diff --git a/Content.Server/DoAfter/DoAfterComponent.cs b/Content.Server/DoAfter/DoAfterComponent.cs
index 383e27160c..75d085b0a6 100644
--- a/Content.Server/DoAfter/DoAfterComponent.cs
+++ b/Content.Server/DoAfter/DoAfterComponent.cs
@@ -11,4 +11,10 @@ namespace Content.Server.DoAfter
// we'll just send them the index. Doesn't matter if it wraps around.
public byte RunningIndex;
}
+
+ ///
+ /// Added to entities that are currently performing any doafters.
+ ///
+ [RegisterComponent]
+ public sealed class ActiveDoAfterComponent : Component {}
}
diff --git a/Content.Server/DoAfter/DoAfterSystem.cs b/Content.Server/DoAfter/DoAfterSystem.cs
index a370d0bad4..3c250a7dec 100644
--- a/Content.Server/DoAfter/DoAfterSystem.cs
+++ b/Content.Server/DoAfter/DoAfterSystem.cs
@@ -26,6 +26,7 @@ namespace Content.Server.DoAfter
public void Add(DoAfterComponent component, DoAfter doAfter)
{
component.DoAfters.Add(doAfter, component.RunningIndex);
+ EnsureComp(component.Owner);
component.RunningIndex++;
Dirty(component);
}
@@ -37,6 +38,11 @@ namespace Content.Server.DoAfter
component.DoAfters.Remove(doAfter);
+ if (component.DoAfters.Count == 0)
+ {
+ RemComp(component.Owner);
+ }
+
RaiseNetworkEvent(new CancelledDoAfterMessage(component.Owner, index));
}
@@ -50,6 +56,11 @@ namespace Content.Server.DoAfter
return;
component.DoAfters.Remove(doAfter);
+
+ if (component.DoAfters.Count == 0)
+ {
+ RemComp(component.Owner);
+ }
}
private void OnDoAfterGetState(EntityUid uid, DoAfterComponent component, ref ComponentGetState args)
@@ -105,7 +116,7 @@ namespace Content.Server.DoAfter
{
base.Update(frameTime);
- foreach (var comp in EntityManager.EntityQuery())
+ foreach (var (_, comp) in EntityManager.EntityQuery())
{
foreach (var (doAfter, _) in comp.DoAfters.ToArray())
{