From 1945c7d7c6ddb1eaa7bbcf81ecb716834910df01 Mon Sep 17 00:00:00 2001
From: SlamBamActionman <83650252+SlamBamActionman@users.noreply.github.com>
Date: Thu, 13 Mar 2025 19:07:52 +0100
Subject: [PATCH] Improvements to antag-before-job selection system (#35822)

* Fix the latejoin-antag-deficit bug, add datafield, add logging

* Fix multiple roles being made for single-role defs,
---
 .../Antag/AntagSelectionSystem.API.cs         | 36 +++++------
 Content.Server/Antag/AntagSelectionSystem.cs  | 59 ++++++++++++++-----
 .../Components/AntagSelectionComponent.cs     |  5 +-
 Content.Shared.Database/LogType.cs            |  7 ++-
 4 files changed, 71 insertions(+), 36 deletions(-)

diff --git a/Content.Server/Antag/AntagSelectionSystem.API.cs b/Content.Server/Antag/AntagSelectionSystem.API.cs
index c89e4df312..93b5fa6136 100644
--- a/Content.Server/Antag/AntagSelectionSystem.API.cs
+++ b/Content.Server/Antag/AntagSelectionSystem.API.cs
@@ -97,7 +97,7 @@ public sealed partial class AntagSelectionSystem
         var countOffset = 0;
         foreach (var otherDef in ent.Comp.Definitions)
         {
-            countOffset += Math.Clamp((poolSize - countOffset) / otherDef.PlayerRatio, otherDef.Min, otherDef.Max) * otherDef.PlayerRatio;
+            countOffset += Math.Clamp((poolSize - countOffset) / otherDef.PlayerRatio, otherDef.Min, otherDef.Max) * otherDef.PlayerRatio; // Note: Is the PlayerRatio necessary here? Seems like it can cause issues for defs with varied PlayerRatio.
         }
         // make sure we don't double-count the current selection
         countOffset -= Math.Clamp(poolSize / def.PlayerRatio, def.Min, def.Max) * def.PlayerRatio;
@@ -362,7 +362,8 @@ public sealed partial class AntagSelectionSystem
     /// <summary>
     /// Get all sessions that have been preselected for antag.
     /// </summary>
-    public HashSet<ICommonSession> GetPreSelectedAntagSessions(AntagSelectionComponent? except = null)
+    /// <param name="except">A specific definition to be excluded from the check.</param>
+    public HashSet<ICommonSession> GetPreSelectedAntagSessions(AntagSelectionDefinition? except = null)
     {
         var result = new HashSet<ICommonSession>();
         var query = QueryAllRules();
@@ -371,15 +372,13 @@ public sealed partial class AntagSelectionSystem
             if (HasComp<EndedGameRuleComponent>(uid))
                 continue;
 
-            if (comp == except)
-                continue;
-
-            if (!comp.PreSelectionsComplete)
-                continue;
-
             foreach (var def in comp.Definitions)
             {
-                result.UnionWith(comp.PreSelectedSessions);
+                if (def.Equals(except))
+                    continue;
+
+                if (comp.PreSelectedSessions.TryGetValue(def, out var set))
+                    result.UnionWith(set);
             }
         }
 
@@ -389,7 +388,11 @@ public sealed partial class AntagSelectionSystem
     /// <summary>
     /// Get all sessions that have been preselected for antag and are exclusive, i.e. should not be paired with other antags.
     /// </summary>
-    public HashSet<ICommonSession> GetPreSelectedExclusiveAntagSessions(AntagSelectionComponent? except = null)
+    /// <param name="except">A specific definition to be excluded from the check.</param>
+    // Note: This is a bit iffy since technically this exclusive definition is defined via the MultiAntagSetting, while there's a separately tracked antagExclusive variable in the mindrole.
+    // We can't query that however since there's no guarantee the mindrole has been given out yet when checking pre-selected antags.
+    // I don't think there's any instance where they differ, but it's something to be aware of for a potential future refactor.
+    public HashSet<ICommonSession> GetPreSelectedExclusiveAntagSessions(AntagSelectionDefinition? except = null)
     {
         var result = new HashSet<ICommonSession>();
         var query = QueryAllRules();
@@ -398,17 +401,14 @@ public sealed partial class AntagSelectionSystem
             if (HasComp<EndedGameRuleComponent>(uid))
                 continue;
 
-            if (comp == except)
-                continue;
-
-            if (!comp.PreSelectionsComplete)
-                continue;
-
             foreach (var def in comp.Definitions)
             {
-                if (def.MultiAntagSetting == AntagAcceptability.None)
+                if (def.Equals(except))
+                    continue;
+
+                if (def.MultiAntagSetting == AntagAcceptability.None && comp.PreSelectedSessions.TryGetValue(def, out var set))
                 {
-                    result.UnionWith(comp.PreSelectedSessions);
+                    result.UnionWith(set);
                     break;
                 }
             }
diff --git a/Content.Server/Antag/AntagSelectionSystem.cs b/Content.Server/Antag/AntagSelectionSystem.cs
index 298fa61a67..c8b2a7d4bb 100644
--- a/Content.Server/Antag/AntagSelectionSystem.cs
+++ b/Content.Server/Antag/AntagSelectionSystem.cs
@@ -12,8 +12,10 @@ using Content.Server.Roles;
 using Content.Server.Roles.Jobs;
 using Content.Server.Shuttles.Components;
 using Content.Server.Station.Events;
+using Content.Shared.Administration.Logs;
 using Content.Shared.Antag;
 using Content.Shared.Clothing;
+using Content.Shared.Database;
 using Content.Shared.GameTicking;
 using Content.Shared.GameTicking.Components;
 using Content.Shared.Ghost;
@@ -47,6 +49,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
     [Dependency] private readonly RoleSystem _role = default!;
     [Dependency] private readonly TransformSystem _transform = default!;
     [Dependency] private readonly EntityWhitelistSystem _whitelist = default!;
+    [Dependency] private readonly ISharedAdminLogManager _adminLogger = default!;
 
     // arbitrary random number to give late joining some mild interest.
     public const float LateJoinRandomChance = 0.5f;
@@ -142,11 +145,13 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
         // eventually this should probably store the players per definition with some kind of unique identifier.
         // something to figure out later.
 
-        var query = QueryActiveRules();
+        var query = QueryAllRules();
         var rules = new List<(EntityUid, AntagSelectionComponent)>();
-        while (query.MoveNext(out var uid, out _, out var antag, out _))
+        while (query.MoveNext(out var uid, out var antag, out _))
         {
-            rules.Add((uid, antag));
+            if (HasComp<ActiveGameRuleComponent>(uid) ||
+                (HasComp<DelayedStartRuleComponent>(uid) && antag.SelectionTime == AntagSelectionTime.IntraPlayerSpawn)) //IntraPlayerSpawn selects antags before spawning, but doesn't activate until after.
+                rules.Add((uid, antag));
         }
         RobustRandom.Shuffle(rules);
 
@@ -158,7 +163,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
             if (!antag.Definitions.Any(p => p.LateJoinAdditional))
                 continue;
 
-            DebugTools.AssertEqual(antag.SelectionTime, AntagSelectionTime.PostPlayerSpawn);
+            DebugTools.AssertNotEqual(antag.SelectionTime, AntagSelectionTime.PrePlayerSpawn);
 
             // do not count players in the lobby for the antag ratio
             var players = _playerManager.NetworkedSessions.Count(x => x.AttachedEntity != null);
@@ -166,7 +171,9 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
             if (!TryGetNextAvailableDefinition((uid, antag), out var def, players))
                 continue;
 
-            if (TryMakeAntag((uid, antag), args.Player, def.Value))
+            var onlyPreSelect = (antag.SelectionTime == AntagSelectionTime.IntraPlayerSpawn && !antag.AssignmentComplete); // Don't wanna give them antag status if the rule hasn't assigned its existing ones yet
+
+            if (TryMakeAntag((uid, antag), args.Player, def.Value, onlyPreSelect: onlyPreSelect))
                 break;
         }
     }
@@ -272,7 +279,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
                     break;
                 }
 
-                if (session != null && ent.Comp.PreSelectedSessions.Contains(session))
+                if (session != null && ent.Comp.PreSelectedSessions.Values.Any(x => x.Contains(session)))
                 {
                     Log.Warning($"Somehow picked {session} for an antag when this rule already selected them previously");
                     continue;
@@ -283,8 +290,11 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
                 MakeAntag(ent, null, def); // This is for spawner antags
             else
             {
-                ent.Comp.PreSelectedSessions.Add(session); // Selection done!
-                Log.Debug($"Selected {session.Name} as antagonist: {ToPrettyString(ent)}");
+                if (!ent.Comp.PreSelectedSessions.TryGetValue(def, out var set))
+                    ent.Comp.PreSelectedSessions.Add(def, set = new HashSet<ICommonSession>());
+                set.Add(session); // Selection done!
+                Log.Debug($"Pre-selected {session.Name} as antagonist: {ToPrettyString(ent)}");
+                _adminLogger.Add(LogType.AntagSelection, $"Pre-selected {session.Name} as antagonist: {ToPrettyString(ent)}");
             }
         }
     }
@@ -300,7 +310,10 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
 
         foreach (var def in ent.Comp.Definitions)
         {
-            foreach (var session in ent.Comp.PreSelectedSessions)
+            if (!ent.Comp.PreSelectedSessions.TryGetValue(def, out var set))
+                continue;
+
+            foreach (var session in set)
             {
                 TryMakeAntag(ent, session, def);
             }
@@ -312,7 +325,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
     /// <summary>
     /// Tries to makes a given player into the specified antagonist.
     /// </summary>
-    public bool TryMakeAntag(Entity<AntagSelectionComponent> ent, ICommonSession? session, AntagSelectionDefinition def, bool ignoreSpawner = false, bool checkPref = true)
+    public bool TryMakeAntag(Entity<AntagSelectionComponent> ent, ICommonSession? session, AntagSelectionDefinition def, bool ignoreSpawner = false, bool checkPref = true, bool onlyPreSelect = false)
     {
         if (checkPref && !HasPrimaryAntagPreference(session, def))
             return false;
@@ -320,7 +333,19 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
         if (!IsSessionValid(ent, session, def) || !IsEntityValid(session?.AttachedEntity, def))
             return false;
 
-        MakeAntag(ent, session, def, ignoreSpawner);
+        if (onlyPreSelect && session != null)
+        {
+            if (!ent.Comp.PreSelectedSessions.TryGetValue(def, out var set))
+                ent.Comp.PreSelectedSessions.Add(def, set = new HashSet<ICommonSession>());
+            set.Add(session);
+            Log.Debug($"Pre-selected {session!.Name} as antagonist: {ToPrettyString(ent)}");
+            _adminLogger.Add(LogType.AntagSelection, $"Pre-selected {session.Name} as antagonist: {ToPrettyString(ent)}");
+        }
+        else
+        {
+            MakeAntag(ent, session, def, ignoreSpawner);
+        }
+
         return true;
     }
 
@@ -334,6 +359,9 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
 
         if (session != null)
         {
+            if (!ent.Comp.PreSelectedSessions.TryGetValue(def, out var set))
+                ent.Comp.PreSelectedSessions.Add(def, set = new HashSet<ICommonSession>());
+            set.Add(session);
             ent.Comp.AssignedSessions.Add(session);
 
             // we shouldn't be blocking the entity if they're just a ghost or smth.
@@ -359,7 +387,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
             if (session != null)
             {
                 ent.Comp.AssignedSessions.Remove(session);
-                ent.Comp.PreSelectedSessions.Remove(session);
+                ent.Comp.PreSelectedSessions[def].Remove(session);
             }
 
             return;
@@ -384,7 +412,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
                 if (session != null)
                 {
                     ent.Comp.AssignedSessions.Remove(session);
-                    ent.Comp.PreSelectedSessions.Remove(session);
+                    ent.Comp.PreSelectedSessions[def].Remove(session);
                 }
 
                 return;
@@ -423,6 +451,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
             SendBriefing(session, def.Briefing);
 
             Log.Debug($"Assigned {ToPrettyString(curMind)} as antagonist: {ToPrettyString(ent)}");
+            _adminLogger.Add(LogType.AntagSelection, $"Assigned {ToPrettyString(curMind)} as antagonist: {ToPrettyString(ent)}");
         }
 
         var afterEv = new AfterAntagEntitySelectedEvent(session, player, ent, def);
@@ -481,7 +510,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
             {
                 if (_role.MindIsAntagonist(mind))
                     return false;
-                if (GetPreSelectedAntagSessions(ent.Comp).Contains(session)) // Used for rules where the antag has been selected, but not started yet
+                if (GetPreSelectedAntagSessions(def).Contains(session)) // Used for rules where the antag has been selected, but not started yet
                     return false;
                 break;
             }
@@ -489,7 +518,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
             {
                 if (_role.MindIsExclusiveAntagonist(mind))
                     return false;
-                if (GetPreSelectedExclusiveAntagSessions(ent.Comp).Contains(session))
+                if (GetPreSelectedExclusiveAntagSessions(def).Contains(session))
                     return false;
                 break;
             }
diff --git a/Content.Server/Antag/Components/AntagSelectionComponent.cs b/Content.Server/Antag/Components/AntagSelectionComponent.cs
index c1e60e8068..5c752ddff7 100644
--- a/Content.Server/Antag/Components/AntagSelectionComponent.cs
+++ b/Content.Server/Antag/Components/AntagSelectionComponent.cs
@@ -44,9 +44,10 @@ public sealed partial class AntagSelectionComponent : Component
     public AntagSelectionTime SelectionTime = AntagSelectionTime.PostPlayerSpawn;
 
     /// <summary>
-    /// Cached sessions of players who are chosen yet not given the role yet.
+    /// Cached sessions of antag definitions and selected players. Players in this dict are not guaranteed to have been assigned the role yet.
     /// </summary>
-    public HashSet<ICommonSession> PreSelectedSessions = new();
+    [DataField]
+    public Dictionary<AntagSelectionDefinition, HashSet<ICommonSession>>PreSelectedSessions = new();
 
     /// <summary>
     /// Cached sessions of players who are chosen. Used so we don't have to rebuild the pool multiple times in a tick.
diff --git a/Content.Shared.Database/LogType.cs b/Content.Shared.Database/LogType.cs
index 5ebb100daf..a868d0e384 100644
--- a/Content.Shared.Database/LogType.cs
+++ b/Content.Shared.Database/LogType.cs
@@ -449,9 +449,14 @@ public enum LogType
     /// An atmos networked device (such as a vent or pump) has had its settings changed, usually through an air alarm
     /// </summary>
     AtmosDeviceSetting = 97,
-    
+
     /// <summary>
     /// Commands related to admemes. Stuff like config changes, etc.
     /// </summary>
     AdminCommands = 98,
+
+    /// <summary>
+    /// A player was selected or assigned antag status
+    /// </summary>
+    AntagSelection = 99,
 }
-- 
2.51.2