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()) {