Improvements and fixups for New Status Effect API (#38660)

Content.Server/Drowsiness/DrowsinessSystem.cs

@@ -1,4 +1,4 @@
-using Content.Server.StatusEffectNew;
+using Content.Server.StatusEffectNew;
 using Content.Shared.Bed.Sleep;
 using Content.Shared.Drowsiness;
 using Content.Shared.StatusEffectNew;
@@ -47,7 +47,7 @@ public sealed class DrowsinessSystem : SharedDrowsinessSystem
             // Make sure the sleep time doesn't cut into the time to next incident.
             drowsiness.NextIncidentTime += duration;
 
-            _statusEffects.TryAddStatusEffect(statusEffect.AppliedTo.Value, SleepingSystem.StatusEffectForcedSleeping, duration);
+            _statusEffects.TryAddStatusEffectDuration(statusEffect.AppliedTo.Value, SleepingSystem.StatusEffectForcedSleeping, duration);
         }
     }
 }

Content.Server/Traits/Assorted/NarcolepsySystem.cs

@@ -53,7 +53,7 @@ public sealed class NarcolepsySystem : EntitySystem
             // Make sure the sleep time doesn't cut into the time to next incident.
             narcolepsy.NextIncidentTime += duration;
 
-            _statusEffects.TryAddStatusEffect(uid, SleepingSystem.StatusEffectForcedSleeping, TimeSpan.FromSeconds(duration));
+            _statusEffects.TryAddStatusEffectDuration(uid, SleepingSystem.StatusEffectForcedSleeping, TimeSpan.FromSeconds(duration));
         }
     }
 }

Content.Shared/EntityEffects/Effects/StatusEffects/ModifyStatusEffect.cs

@@ -20,7 +20,7 @@ public sealed partial class ModifyStatusEffect : EntityEffect
     public float Time = 2.0f;
 
     /// <remarks>
-    /// true - refresh status effect time, false - accumulate status effect time.
+    /// true - refresh status effect time (update to greater value), false - accumulate status effect time.
     /// </remarks>
     [DataField]
     public bool Refresh = true;
@@ -40,16 +40,20 @@ public sealed partial class ModifyStatusEffect : EntityEffect
         if (args is EntityEffectReagentArgs reagentArgs)
             time *= reagentArgs.Scale.Float();
 
+        var duration = TimeSpan.FromSeconds(time);
         switch (Type)
         {
             case StatusEffectMetabolismType.Add:
-                statusSys.TryAddStatusEffect(args.TargetEntity, EffectProto, TimeSpan.FromSeconds(time), Refresh);
+                if (Refresh)
+                    statusSys.TryUpdateStatusEffectDuration(args.TargetEntity, EffectProto, duration);
+                else
+                    statusSys.TryAddStatusEffectDuration(args.TargetEntity, EffectProto, duration);
                 break;
             case StatusEffectMetabolismType.Remove:
-                statusSys.TryAddTime(args.TargetEntity, EffectProto, -TimeSpan.FromSeconds(time));
+                statusSys.TryAddTime(args.TargetEntity, EffectProto, -duration);
                 break;
             case StatusEffectMetabolismType.Set:
-                statusSys.TrySetTime(args.TargetEntity, EffectProto, TimeSpan.FromSeconds(time));
+                statusSys.TrySetStatusEffectDuration(args.TargetEntity, EffectProto, duration);
                 break;
         }
     }

Content.Shared/SSDIndicator/SSDIndicatorSystem.cs

@@ -85,7 +85,7 @@ public sealed class SSDIndicatorSystem : EntitySystem
                 ssd.FallAsleepTime <= _timing.CurTime &&
                 !TerminatingOrDeleted(uid))
             {
-                _statusEffects.TryAddStatusEffect(uid, StatusEffectSSDSleeping);
+                _statusEffects.TrySetStatusEffectDuration(uid, StatusEffectSSDSleeping, null);
             }
         }
     }

Content.Shared/StatusEffectNew/SharedStatusEffectsSystem.cs

@@ -1,3 +1,4 @@
+using System.Diagnostics.CodeAnalysis;
 using Content.Shared.Alert;
 using Content.Shared.StatusEffectNew.Components;
 using Content.Shared.Whitelist;
@@ -77,53 +78,59 @@ public abstract partial class SharedStatusEffectsSystem : EntitySystem
         effectComp.EndEffectTime += delta;
         Dirty(effect, effectComp);
 
-        if (effectComp is { AppliedTo: not null, Alert: not null })
+        ShowAlertIfNeeded(effectComp);
-        {
-            (TimeSpan Start, TimeSpan End)? cooldown = effectComp.EndEffectTime is null
-                ? null
-                : (_timing.CurTime, effectComp.EndEffectTime.Value);
-            _alerts.ShowAlert(
-                effectComp.AppliedTo.Value,
-                effectComp.Alert.Value,
-                cooldown: cooldown
-            );
-        }
     }
 
-    private void SetStatusEffectTime(EntityUid effect, TimeSpan duration)
+    private void SetStatusEffectTime(EntityUid effect, TimeSpan? duration)
     {
         if (!_effectQuery.TryComp(effect, out var effectComp))
             return;
 
-        effectComp.EndEffectTime = _timing.CurTime + duration;
+        if (duration is null)
+        {
+            if(effectComp.EndEffectTime is null)
+                return;
+
+            effectComp.EndEffectTime = null;
+        }
+        else
+            effectComp.EndEffectTime = _timing.CurTime + duration;
+
         Dirty(effect, effectComp);
 
-        if (effectComp is { AppliedTo: not null, Alert: not null })
+        ShowAlertIfNeeded(effectComp);
-        {
-            (TimeSpan, TimeSpan)? cooldown = effectComp.EndEffectTime is null
-                ? null
-                : (_timing.CurTime, effectComp.EndEffectTime.Value);
-            _alerts.ShowAlert(
-                effectComp.AppliedTo.Value,
-                effectComp.Alert.Value,
-                cooldown: cooldown
-            );
-        }
     }
 
+    private void UpdateStatusEffectTime(EntityUid effect, TimeSpan? duration)
+    {
+        if (!_effectQuery.TryComp(effect, out var effectComp))
+            return;
+
+        // It's already infinitely long
+        if (effectComp.EndEffectTime is null)
+            return;
+
+        if (duration is null)
+            effectComp.EndEffectTime = null;
+        else
+        {
+            var newEndTime = _timing.CurTime + duration;
+            if (effectComp.EndEffectTime >= newEndTime)
+                return;
+
+            effectComp.EndEffectTime = newEndTime;
+        }
+
+        Dirty(effect, effectComp);
+
+        ShowAlertIfNeeded(effectComp);
+    }
+
+
     private void OnStatusEffectApplied(Entity<StatusEffectComponent> ent, ref StatusEffectAppliedEvent args)
     {
-        if (ent.Comp is { AppliedTo: not null, Alert: not null })
+        StatusEffectComponent statusEffect = ent;
-        {
+        ShowAlertIfNeeded(statusEffect);
-            (TimeSpan, TimeSpan)? cooldown = ent.Comp.EndEffectTime is null
-                ? null
-                : (_timing.CurTime, ent.Comp.EndEffectTime.Value);
-            _alerts.ShowAlert(
-                ent.Comp.AppliedTo.Value,
-                ent.Comp.Alert.Value,
-                cooldown: cooldown
-            );
-        }
     }
 
     private void OnStatusEffectRemoved(Entity<StatusEffectComponent> ent, ref StatusEffectRemovedEvent args)
@@ -154,6 +161,64 @@ public abstract partial class SharedStatusEffectsSystem : EntitySystem
 
         return true;
     }
+
+    /// <summary>
+    /// Attempts to add a status effect to the specified entity. Returns True if the effect is added, does not check if one
+    /// already exists as it's intended to be called after a check for an existing effect has already failed.
+    /// </summary>
+    /// <param name="target">The target entity to which the effect should be added.</param>
+    /// <param name="effectProto">ProtoId of the status effect entity. Make sure it has StatusEffectComponent on it.</param>
+    /// <param name="duration">Duration of status effect. Leave null and the effect will be permanent until it is removed using <c>TryRemoveStatusEffect</c>.</param>
+    /// <param name="statusEffect">The EntityUid of the status effect we have just created or null if we couldn't create one.</param>
+    private bool TryAddStatusEffect(
+        EntityUid target,
+        EntProtoId effectProto,
+        [NotNullWhen(true)] out EntityUid? statusEffect,
+        TimeSpan? duration = null
+    )
+    {
+        statusEffect = null;
+        if (!CanAddStatusEffect(target, effectProto))
+            return false;
+
+        var container = EnsureComp<StatusEffectContainerComponent>(target);
+
+        //And only if all checks passed we spawn the effect
+        var effect = PredictedSpawnAttachedTo(effectProto, Transform(target).Coordinates);
+        _transform.SetParent(effect, target);
+        if (!_effectQuery.TryComp(effect, out var effectComp))
+            return false;
+
+        statusEffect = effect;
+
+        if (duration != null)
+            effectComp.EndEffectTime = _timing.CurTime + duration;
+
+        container.ActiveStatusEffects.Add(effect);
+        effectComp.AppliedTo = target;
+        Dirty(target, container);
+        Dirty(effect, effectComp);
+
+        var ev = new StatusEffectAppliedEvent(target);
+        RaiseLocalEvent(effect, ref ev);
+
+        return true;
+    }
+
+    private void ShowAlertIfNeeded(StatusEffectComponent effectComp)
+    {
+        if (effectComp is { AppliedTo: not null, Alert: not null })
+        {
+            (TimeSpan, TimeSpan)? cooldown = effectComp.EndEffectTime is null
+                ? null
+                : (_timing.CurTime, effectComp.EndEffectTime.Value);
+            _alerts.ShowAlert(
+                effectComp.AppliedTo.Value,
+                effectComp.Alert.Value,
+                cooldown: cooldown
+            );
+        }
+    }
 }
 
 /// <summary>

Content.Shared/StatusEffectNew/StatusEffectSystem.API.cs

@@ -7,80 +7,94 @@ namespace Content.Shared.StatusEffectNew;
 public abstract partial class SharedStatusEffectsSystem
 {
     /// <summary>
-    /// Attempts to add a status effect to the specified entity. Returns True if the effect is added or it already exists
+    /// Increments duration of status effect by <see cref="duration"/>.
-    /// and has been successfully extended in time, returns False if the status effect cannot be applied to this entity,
+    /// Tries to add status effect if it is not yet present on entity.
-    /// or for any other reason.
     /// </summary>
     /// <param name="target">The target entity to which the effect should be added.</param>
     /// <param name="effectProto">ProtoId of the status effect entity. Make sure it has StatusEffectComponent on it.</param>
     /// <param name="duration">Duration of status effect. Leave null and the effect will be permanent until it is removed using <c>TryRemoveStatusEffect</c>.</param>
-    /// <param name="resetCooldown">
-    /// If True, the effect duration time will be reset and reapplied. If False, the effect duration time will be overlaid with the existing one.
-    /// In the other case, the effect will either be added for the specified time or its time will be extended for the specified time.
-    /// </param>
     /// <param name="statusEffect">The EntityUid of the status effect we have just created or null if it doesn't exist.</param>
-    public bool TryAddStatusEffect(
+    /// <returns>True if effect exists and its duration is set properly, false in case effect cannot be applied.</returns>
+    public bool TryAddStatusEffectDuration(
         EntityUid target,
         EntProtoId effectProto,
-        out EntityUid? statusEffect,
+        [NotNullWhen(true)] out EntityUid? statusEffect,
-        TimeSpan? duration = null,
+        TimeSpan duration
-        bool resetCooldown = false
     )
     {
-        statusEffect = null;
+        if (!TryGetStatusEffect(target, effectProto, out statusEffect))
-        if (TryGetStatusEffect(target, effectProto, out var existingEffect))
+            return TryAddStatusEffect(target, effectProto, out statusEffect, duration);
-        {
-            statusEffect = existingEffect;
-            //We don't need to add the effect if it already exists
-            if (duration is null)
-                return true;
 
-            if (resetCooldown)
+        AddStatusEffectTime(statusEffect.Value, duration);
-                SetStatusEffectTime(existingEffect.Value, duration.Value);
-            else
-                AddStatusEffectTime(existingEffect.Value, duration.Value);
-
-            return true;
-        }
-
-        if (!CanAddStatusEffect(target, effectProto))
-            return false;
-
-        var container = EnsureComp<StatusEffectContainerComponent>(target);
-
-        //And only if all checks passed we spawn the effect
-        var effect = PredictedSpawnAttachedTo(effectProto, Transform(target).Coordinates);
-        statusEffect = effect;
-        _transform.SetParent(effect, target);
-        if (!_effectQuery.TryComp(effect, out var effectComp))
-            return false;
-
-        if (duration != null)
-            effectComp.EndEffectTime = _timing.CurTime + duration;
-
-        container.ActiveStatusEffects.Add(effect);
-        effectComp.AppliedTo = target;
-        Dirty(target, container);
-        Dirty(effect, effectComp);
-
-        var ev = new StatusEffectAppliedEvent(target);
-        RaiseLocalEvent(effect, ref ev);
 
         return true;
     }
 
+
+    ///<inheritdoc cref="TryAddStatusEffectDuration(Robust.Shared.GameObjects.EntityUid,Robust.Shared.Prototypes.EntProtoId,out Robust.Shared.GameObjects.EntityUid?,System.TimeSpan)"/>
+    public bool TryAddStatusEffectDuration(EntityUid target, EntProtoId effectProto, TimeSpan duration)
+    {
+        return TryAddStatusEffectDuration(target, effectProto, out _, duration);
+    }
+
     /// <summary>
-    /// An overload of <see cref="TryAddStatusEffect(EntityUid,EntProtoId,out EntityUid?,TimeSpan?,bool)"/>
+    /// Sets duration of status effect by <see cref="duration"/>.
-    /// that doesn't return a status effect EntityUid.
+    /// Tries to add status effect if it is not yet present on entity.
     /// </summary>
-    public bool TryAddStatusEffect(
+    /// <param name="target">The target entity to which the effect should be added.</param>
+    /// <param name="effectProto">ProtoId of the status effect entity. Make sure it has StatusEffectComponent on it.</param>
+    /// <param name="duration">Duration of status effect. Leave null and the effect will be permanent until it is removed using <c>TryRemoveStatusEffect</c>.</param>
+    /// <param name="statusEffect">The EntityUid of the status effect we have just created or null if it doesn't exist.</param>
+    /// <returns>True if effect exists and its duration is set properly, false in case effect cannot be applied.</returns>
+    public bool TrySetStatusEffectDuration(
         EntityUid target,
         EntProtoId effectProto,
-        TimeSpan? duration = null,
+        [NotNullWhen(true)] out EntityUid? statusEffect,
-        bool resetCooldown = false
+        TimeSpan? duration = null
     )
     {
-        return TryAddStatusEffect(target, effectProto, out _, duration, resetCooldown);
+        if (!TryGetStatusEffect(target, effectProto, out statusEffect))
+            return TryAddStatusEffect(target, effectProto, out statusEffect, duration);
+
+        SetStatusEffectTime(statusEffect.Value, duration);
+
+        return true;
+    }
+
+    /// <inheritdoc cref="TrySetStatusEffectDuration(Robust.Shared.GameObjects.EntityUid,Robust.Shared.Prototypes.EntProtoId,out Robust.Shared.GameObjects.EntityUid?,System.TimeSpan?)"/>
+    public bool TrySetStatusEffectDuration(EntityUid target, EntProtoId effectProto, TimeSpan? duration = null)
+    {
+        return TrySetStatusEffectDuration(target, effectProto, out _, duration);
+    }
+
+    /// <summary>
+    /// Updates duration of effect to larger value between provided <see cref="duration"/> and current effect duration.
+    /// Tries to add status effect if it is not yet present on entity.
+    /// </summary>
+    /// <param name="target">The target entity to which the effect should be added.</param>
+    /// <param name="effectProto">ProtoId of the status effect entity. Make sure it has StatusEffectComponent on it.</param>
+    /// <param name="duration">Duration of status effect. Leave null and the effect will be permanent until it is removed using <c>TryRemoveStatusEffect</c>.</param>
+    /// <param name="statusEffect">The EntityUid of the status effect we have just created or null if it doesn't exist.</param>
+    /// <returns>True if effect exists and its duration is set properly, false in case effect cannot be applied.</returns>
+    public bool TryUpdateStatusEffectDuration(
+        EntityUid target,
+        EntProtoId effectProto,
+        [NotNullWhen(true)] out EntityUid? statusEffect,
+        TimeSpan? duration = null
+    )
+    {
+        if (!TryGetStatusEffect(target, effectProto, out statusEffect))
+            return TryAddStatusEffect(target, effectProto, out statusEffect, duration);
+
+        UpdateStatusEffectTime(statusEffect.Value, duration);
+
+        return true;
+    }
+
+    /// <inheritdoc cref="TryUpdateStatusEffectDuration(Robust.Shared.GameObjects.EntityUid,Robust.Shared.Prototypes.EntProtoId,out Robust.Shared.GameObjects.EntityUid?,System.TimeSpan?)"/>
+    public bool TryUpdateStatusEffectDuration(EntityUid target, EntProtoId effectProto, TimeSpan? duration = null)
+    {
+        return TryUpdateStatusEffectDuration(target, effectProto, out _, duration);
     }
 
     /// <summary>
@@ -190,6 +204,36 @@ public abstract partial class SharedStatusEffectsSystem
         return false;
     }
 
+    /// <summary>
+    /// Attempts to get the maximum time left for a given Status Effect Component, returns false if no such
+    /// component exists.
+    /// </summary>
+    /// <param name="uid">The target entity on which the effect is applied.</param>
+    /// <param name="time">Returns the EntityUid of the status effect with the most time left, and the end effect time
+    /// of that status effect.</param>
+    /// <returns> True if a status effect entity with the given component exists</returns>
+    public bool TryGetMaxTime<T>(EntityUid uid, out (EntityUid EffectEnt, TimeSpan? EndEffectTime) time) where T : IComponent
+    {
+        time = default;
+        if (!TryEffectsWithComp<T>(uid, out var status))
+            return false;
+
+        time.Item2 = TimeSpan.Zero;
+
+        foreach (var effect in status)
+        {
+            if (effect.Comp2.EndEffectTime == null)
+            {
+                time = (effect.Owner, null);
+                return true;
+            }
+
+            if (effect.Comp2.EndEffectTime > time.Item2)
+                time = (effect.Owner, effect.Comp2.EndEffectTime);
+        }
+        return true;
+    }
+
     /// <summary>
     /// Attempts to edit the remaining time for a status effect on an entity.
     /// </summary>