Enhances Roundstart Antagonist Selection to Prevent Roundstart Fail and Allocate All Antags

Content.Server/Antag/AntagSelectionPlayerPool.cs

@@ -2,6 +2,7 @@
 using System.Linq;
 using Robust.Shared.Player;
 using Robust.Shared.Random;
+using Robust.Shared.Utility;
 
 namespace Content.Server.Antag;
 
@@ -23,5 +24,46 @@ public bool TryPickAndTake(IRobustRandom random, [NotNullWhen(true)] out ICommon
         return session != null;
     }
 
+    // EE
+    public bool TryGetItems(IRobustRandom random,
+                            [NotNullWhen(true)] out ICommonSession[]? sessions,
+                            int count,
+                            bool allowDuplicates = true)
+    {
+        DebugTools.Assert(count > 0, $"The count {nameof(count)} of requested sessions must be greater than zero!");
+
+        sessions = null;
+        List<ICommonSession> session_list = [];
+
+        foreach (var pool in orderedPools)
+        {
+            if (pool.Count == 0)
+                continue;
+
+            var picked = random.GetItems(pool, count - session_list.Count, allowDuplicates);
+            session_list.AddRange(picked);
+            if (session_list.Count < count)
+            {
+                continue;
+            }
+            sessions = session_list.ToArray();
+            break;
+        }
+
+        return sessions != null;
+    }
+
+    // EE
+    public AntagSelectionPlayerPool Where(Func<ICommonSession, bool> predicate)
+    {
+        var newPools = orderedPools.Select(
+            (pool) =>
+            {
+                return pool.Where(predicate).ToList();
+            }
+        );
+
+        return new AntagSelectionPlayerPool(newPools.ToList());
+    }
     public int Count => orderedPools.Sum(p => p.Count);
 }

Content.Server/Antag/AntagSelectionSystem.cs

@@ -204,20 +204,68 @@ public void ChooseAntags(Entity<AntagSelectionComponent> ent, IList<ICommonSessi
         var playerPool = GetPlayerPool(ent, pool, def);
         var count = GetTargetAntagCount(ent, GetTotalPlayerCount(pool), def);
 
-        for (var i = 0; i < count; i++)
+        ///// Einstein Engines changes /////
+        //
+        // Fixes issues caused by failures in making someone antag while
+        // not breaking any API on this system.
+        //
+        // This will either allocate `count` slots from the player pool,
+        // or will call MakeAntag with null sessions to fill up the slots.
+        //
+        if (def.PickPlayer)
         {
-            var session = (ICommonSession?) null;
-            if (def.PickPlayer)
+            // Tries multiple times to assign antags.
+            // When any number of assignments fail, next iteration
+            // gets new items to replace those.
+            // Already selected or failed sessions are avoided.
+            // It retries until it ends with no failures or up
+            // to maxRetries attempts.
+
+            const int maxRetries = 4;
+            var retry = 0;
+            List<ICommonSession> failed = [];
+
+            while (ent.Comp.SelectedSessions.Count < count && retry < maxRetries)
             {
-                if (!playerPool.TryPickAndTake(RobustRandom, out session))
+                var sessions = (ICommonSession[]?) null;
+                if (!playerPool.TryGetItems(RobustRandom,
+                                            out sessions,
+                                            count - ent.Comp.SelectedSessions.Count,
+                                            false))
+                    break; // Ends early if there are no eligible sessions
+
+                foreach (var session in sessions)
+                {
+                    MakeAntag(ent, session, def);
+                    if (!ent.Comp.SelectedSessions.Contains(session))
+                    {
+                        failed.Add(session);
+                    }
+                }
+                // In case we're done
+                if (ent.Comp.SelectedSessions.Count >= count)
                     break;
 
-                if (ent.Comp.SelectedSessions.Contains(session))
-                    continue;
+                playerPool = playerPool.Where((session_) =>
+                {
+                    return !ent.Comp.SelectedSessions.Contains(session_) &&
+                        !failed.Contains(session_);
+                });
+                retry++;
             }
+        }
+
+        // This preserves previous behavior for when def.PickPlayer
+        // was not satisfied. This behavior is not that obvious to
+        // read from the previous code.
+        // It may otherwise process leftover slots if maxRetries have
+        // been reached.
 
-            MakeAntag(ent, session, def);
+        for (var i = ent.Comp.SelectedSessions.Count; i < count; i++)
+        {
+            MakeAntag(ent, null, def);
         }
+        ///// End of Einstein Engines changes /////
     }
 
     /// <summary>
@@ -261,17 +309,16 @@ public void MakeAntag(Entity<AntagSelectionComponent> ent, ICommonSession? sessi
         {
             var getEntEv = new AntagSelectEntityEvent(session, ent);
             RaiseLocalEvent(ent, ref getEntEv, true);
-
-            if (!getEntEv.Handled)
-            {
-                throw new InvalidOperationException($"Attempted to make {session} antagonist in gamerule {ToPrettyString(ent)} but there was no valid entity for player.");
-            }
-
             antagEnt = getEntEv.Entity;
         }
 
         if (antagEnt is not { } player)
+        {
+            Log.Error($"Attempted to make {session} antagonist in gamerule {ToPrettyString(ent)} but there was no valid entity for player.");
+            if (session != null)
+                ent.Comp.SelectedSessions.Remove(session);
             return;
+        }
 
         var getPosEv = new AntagSelectLocationEvent(session, ent);
         RaiseLocalEvent(ent, ref getPosEv, true);
@@ -282,11 +329,15 @@ public void MakeAntag(Entity<AntagSelectionComponent> ent, ICommonSession? sessi
             _transform.SetMapCoordinates((player, playerXform), pos);
         }
 
+        // If we want to just do a ghost role spawner, set up data here and then return early.
+        // This could probably be an event in the future if we want to be more refined about it.
         if (isSpawner)
         {
             if (!TryComp<GhostRoleAntagSpawnerComponent>(player, out var spawnerComp))
             {
-                Log.Error("Antag spawner with GhostRoleAntagSpawnerComponent.");
+                Log.Error($"Antag spawner {player} does not have a GhostRoleAntagSpawnerComponent.");
+                if (session != null)
+                    ent.Comp.SelectedSessions.Remove(session);
                 return;
             }