Add support for antag-before-job selection (#35789)

* Add support for antag-before-job selection

* Include logging

Content.Server/Antag/AntagSelectionSystem.API.cs

@@ -3,7 +3,9 @@ using System.Linq;
 using Content.Server.Antag.Components;
 using Content.Server.GameTicking.Rules.Components;
 using Content.Server.Objectives;
+using Content.Shared.Antag;
 using Content.Shared.Chat;
+using Content.Shared.GameTicking.Components;
 using Content.Shared.Mind;
 using Content.Shared.Preferences;
 using JetBrains.Annotations;
@@ -25,7 +27,7 @@ public sealed partial class AntagSelectionSystem
         definition = null;
 
         var totalTargetCount = GetTargetAntagCount(ent, players);
-        var mindCount = ent.Comp.SelectedMinds.Count;
+        var mindCount = ent.Comp.AssignedMinds.Count;
         if (mindCount >= totalTargetCount)
             return false;
 
@@ -115,7 +117,7 @@ public sealed partial class AntagSelectionSystem
             return new List<(EntityUid, SessionData, string)>();
 
         var output = new List<(EntityUid, SessionData, string)>();
-        foreach (var (mind, name) in ent.Comp.SelectedMinds)
+        foreach (var (mind, name) in ent.Comp.AssignedMinds)
         {
             if (!TryComp<MindComponent>(mind, out var mindComp) || mindComp.OriginalOwnerUserId == null)
                 continue;
@@ -137,7 +139,7 @@ public sealed partial class AntagSelectionSystem
             return new();
 
         var output = new List<Entity<MindComponent>>();
-        foreach (var (mind, _) in ent.Comp.SelectedMinds)
+        foreach (var (mind, _) in ent.Comp.AssignedMinds)
         {
             if (!TryComp<MindComponent>(mind, out var mindComp) || mindComp.OriginalOwnerUserId == null)
                 continue;
@@ -155,7 +157,7 @@ public sealed partial class AntagSelectionSystem
         if (!Resolve(ent, ref ent.Comp, false))
             return new();
 
-        return ent.Comp.SelectedMinds.Select(p => p.Item1).ToList();
+        return ent.Comp.AssignedMinds.Select(p => p.Item1).ToList();
     }
 
     /// <summary>
@@ -247,7 +249,7 @@ public sealed partial class AntagSelectionSystem
         if (!Resolve(ent, ref ent.Comp, false))
             return false;
 
-        return GetAliveAntagCount(ent) == ent.Comp.SelectedMinds.Count;
+        return GetAliveAntagCount(ent) == ent.Comp.AssignedMinds.Count;
     }
 
     /// <summary>
@@ -352,8 +354,66 @@ public sealed partial class AntagSelectionSystem
         var ruleEnt = GameTicker.AddGameRule(id);
         RemComp<LoadMapRuleComponent>(ruleEnt);
         var antag = Comp<AntagSelectionComponent>(ruleEnt);
-        antag.SelectionsComplete = true; // don't do normal selection.
+        antag.AssignmentComplete = true; // don't do normal selection.
         GameTicker.StartGameRule(ruleEnt);
         return (ruleEnt, antag);
     }
+
+    /// <summary>
+    /// Get all sessions that have been preselected for antag.
+    /// </summary>
+    public HashSet<ICommonSession> GetPreSelectedAntagSessions(AntagSelectionComponent? except = null)
+    {
+        var result = new HashSet<ICommonSession>();
+        var query = QueryAllRules();
+        while (query.MoveNext(out var uid, out var comp, out _))
+        {
+            if (HasComp<EndedGameRuleComponent>(uid))
+                continue;
+
+            if (comp == except)
+                continue;
+
+            if (!comp.PreSelectionsComplete)
+                continue;
+
+            foreach (var def in comp.Definitions)
+            {
+                result.UnionWith(comp.PreSelectedSessions);
+            }
+        }
+
+        return result;
+    }
+
+    /// <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)
+    {
+        var result = new HashSet<ICommonSession>();
+        var query = QueryAllRules();
+        while (query.MoveNext(out var uid, out var comp, out _))
+        {
+            if (HasComp<EndedGameRuleComponent>(uid))
+                continue;
+
+            if (comp == except)
+                continue;
+
+            if (!comp.PreSelectionsComplete)
+                continue;
+
+            foreach (var def in comp.Definitions)
+            {
+                if (def.MultiAntagSetting == AntagAcceptability.None)
+                {
+                    result.UnionWith(comp.PreSelectedSessions);
+                    break;
+                }
+            }
+        }
+
+        return result;
+    }
 }

Content.Server/Antag/AntagSelectionSystem.cs

@@ -11,6 +11,7 @@ using Content.Server.Preferences.Managers;
 using Content.Server.Roles;
 using Content.Server.Roles.Jobs;
 using Content.Server.Shuttles.Components;
+using Content.Server.Station.Events;
 using Content.Shared.Antag;
 using Content.Shared.Clothing;
 using Content.Shared.GameTicking;
@@ -89,19 +90,33 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
         var query = QueryActiveRules();
         while (query.MoveNext(out var uid, out _, out var comp, out _))
         {
-            if (comp.SelectionTime != AntagSelectionTime.PrePlayerSpawn)
+            if (comp.SelectionTime != AntagSelectionTime.PrePlayerSpawn && comp.SelectionTime != AntagSelectionTime.IntraPlayerSpawn)
                 continue;
 
-            if (comp.SelectionsComplete)
+            if (comp.AssignmentComplete)
+                continue;
+
+            ChooseAntags((uid, comp), pool); // We choose the antags here...
+
+            if (comp.SelectionTime == AntagSelectionTime.PrePlayerSpawn)
+            {
+                AssignPreSelectedSessions((uid, comp)); // ...But only assign them if PrePlayerSpawn
+                foreach (var session in comp.AssignedSessions)
+                {
+                    args.PlayerPool.Remove(session);
+                    GameTicker.PlayerJoinGame(session);
+                }
+            }
+        }
+
+        // If IntraPlayerSpawn is selected, delayed rules should choose at this point too.
+        var queryDelayed = QueryDelayedRules();
+        while (queryDelayed.MoveNext(out var uid, out _, out var comp, out _))
+        {
+            if (comp.SelectionTime != AntagSelectionTime.IntraPlayerSpawn)
                 continue;
 
             ChooseAntags((uid, comp), pool);
-
-            foreach (var session in comp.SelectedSessions)
-            {
-                args.PlayerPool.Remove(session);
-                GameTicker.PlayerJoinGame(session);
-            }
         }
     }
 
@@ -110,10 +125,11 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
         var query = QueryActiveRules();
         while (query.MoveNext(out var uid, out _, out var comp, out _))
         {
-            if (comp.SelectionTime != AntagSelectionTime.PostPlayerSpawn)
+            if (comp.SelectionTime != AntagSelectionTime.PostPlayerSpawn && comp.SelectionTime != AntagSelectionTime.IntraPlayerSpawn)
                 continue;
 
             ChooseAntags((uid, comp), args.Players);
+            AssignPreSelectedSessions((uid, comp));
         }
     }
 
@@ -183,14 +199,20 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
         if (GameTicker.RunLevel != GameRunLevel.InRound)
             return;
 
-        if (component.SelectionsComplete)
+        if (component.AssignmentComplete)
             return;
 
-        var players = _playerManager.Sessions
+        if (!component.PreSelectionsComplete)
-            .Where(x => GameTicker.PlayerGameStatuses.TryGetValue(x.UserId, out var status) && status == PlayerGameStatus.JoinedGame)
+        {
-            .ToList();
+            var players = _playerManager.Sessions
+                .Where(x => GameTicker.PlayerGameStatuses.TryGetValue(x.UserId, out var status) &&
+                            status == PlayerGameStatus.JoinedGame)
+                .ToList();
 
-        ChooseAntags((uid, component), players, midround: true);
+            ChooseAntags((uid, component), players, midround: true);
+        }
+
+        AssignPreSelectedSessions((uid, component));
     }
 
     /// <summary>
@@ -201,7 +223,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
     /// <param name="midround">Disable picking players for pre-spawn antags in the middle of a round</param>
     public void ChooseAntags(Entity<AntagSelectionComponent> ent, IList<ICommonSession> pool, bool midround = false)
     {
-        if (ent.Comp.SelectionsComplete)
+        if (ent.Comp.PreSelectionsComplete)
             return;
 
         foreach (var def in ent.Comp.Definitions)
@@ -209,7 +231,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
             ChooseAntags(ent, pool, def, midround: midround);
         }
 
-        ent.Comp.SelectionsComplete = true;
+        ent.Comp.PreSelectionsComplete = true;
     }
 
     /// <summary>
@@ -250,17 +272,43 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
                     break;
                 }
 
-                if (session != null && ent.Comp.SelectedSessions.Contains(session))
+                if (session != null && ent.Comp.PreSelectedSessions.Contains(session))
                 {
                     Log.Warning($"Somehow picked {session} for an antag when this rule already selected them previously");
                     continue;
                 }
             }
 
-            MakeAntag(ent, session, def);
+            if (session == null)
+                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)}");
+            }
         }
     }
 
+    /// <summary>
+    /// Assigns antag roles to sessions selected for it.
+    /// </summary>
+    public void AssignPreSelectedSessions(Entity<AntagSelectionComponent> ent)
+    {
+        // Only assign if there's been a pre-selection, and the selection hasn't already been made
+        if (!ent.Comp.PreSelectionsComplete || ent.Comp.AssignmentComplete)
+            return;
+
+        foreach (var def in ent.Comp.Definitions)
+        {
+            foreach (var session in ent.Comp.PreSelectedSessions)
+            {
+                TryMakeAntag(ent, session, def);
+            }
+        }
+
+        ent.Comp.AssignmentComplete = true;
+    }
+
     /// <summary>
     /// Tries to makes a given player into the specified antagonist.
     /// </summary>
@@ -286,7 +334,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
 
         if (session != null)
         {
-            ent.Comp.SelectedSessions.Add(session);
+            ent.Comp.AssignedSessions.Add(session);
 
             // we shouldn't be blocking the entity if they're just a ghost or smth.
             if (!HasComp<GhostComponent>(session.AttachedEntity))
@@ -309,7 +357,11 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
         {
             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);
+            {
+                ent.Comp.AssignedSessions.Remove(session);
+                ent.Comp.PreSelectedSessions.Remove(session);
+            }
+
             return;
         }
 
@@ -330,7 +382,11 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
             {
                 Log.Error($"Antag spawner {player} does not have a GhostRoleAntagSpawnerComponent.");
                 if (session != null)
-                    ent.Comp.SelectedSessions.Remove(session);
+                {
+                    ent.Comp.AssignedSessions.Remove(session);
+                    ent.Comp.PreSelectedSessions.Remove(session);
+                }
+
                 return;
             }
 
@@ -363,10 +419,10 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
 
             _mind.TransferTo(curMind.Value, antagEnt, ghostCheckOverride: true);
             _role.MindAddRoles(curMind.Value, def.MindRoles, null, true);
-            ent.Comp.SelectedMinds.Add((curMind.Value, Name(player)));
+            ent.Comp.AssignedMinds.Add((curMind.Value, Name(player)));
             SendBriefing(session, def.Briefing);
 
-            Log.Debug($"Selected {ToPrettyString(curMind)} as antagonist: {ToPrettyString(ent)}");
+            Log.Debug($"Assigned {ToPrettyString(curMind)} as antagonist: {ToPrettyString(ent)}");
         }
 
         var afterEv = new AfterAntagEntitySelectedEvent(session, player, ent, def);
@@ -412,15 +468,11 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
         if (session.Status is SessionStatus.Disconnected or SessionStatus.Zombie)
             return false;
 
-        if (ent.Comp.SelectedSessions.Contains(session))
+        if (ent.Comp.AssignedSessions.Contains(session))
             return false;
 
         mind ??= session.GetMind();
 
-        // If the player has not spawned in as any entity (e.g., in the lobby), they can be given an antag role/entity.
-        if (mind == null)
-            return true;
-
         //todo: we need some way to check that we're not getting the same role twice. (double picking thieves or zombies through midrounds)
 
         switch (def.MultiAntagSetting)
@@ -429,12 +481,16 @@ 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
+                    return false;
                 break;
             }
             case AntagAcceptability.NotExclusive:
             {
                 if (_role.MindIsExclusiveAntagonist(mind))
                     return false;
+                if (GetPreSelectedExclusiveAntagSessions(ent.Comp).Contains(session))
+                    return false;
                 break;
             }
         }
@@ -481,7 +537,7 @@ public sealed partial class AntagSelectionSystem : GameRuleSystem<AntagSelection
         if (ent.Comp.AgentName is not { } name)
             return;
 
-        args.Minds = ent.Comp.SelectedMinds;
+        args.Minds = ent.Comp.AssignedMinds;
         args.AgentName = Loc.GetString(name);
     }
 }

Content.Server/Antag/Components/AntagSelectionComponent.cs

@@ -14,10 +14,16 @@ namespace Content.Server.Antag.Components;
 public sealed partial class AntagSelectionComponent : Component
 {
     /// <summary>
-    /// Has the primary selection of antagonists finished yet?
+    /// Has the primary assignment of antagonists finished yet?
     /// </summary>
     [DataField]
-    public bool SelectionsComplete;
+    public bool AssignmentComplete;
+
+    /// <summary>
+    /// Has the antagonists been preselected but yet to be fully assigned?
+    /// </summary>
+    [DataField]
+    public bool PreSelectionsComplete;
 
     /// <summary>
     /// The definitions for the antagonists
@@ -26,10 +32,10 @@ public sealed partial class AntagSelectionComponent : Component
     public List<AntagSelectionDefinition> Definitions = new();
 
     /// <summary>
-    /// The minds and original names of the players selected to be antagonists.
+    /// The minds and original names of the players assigned to be antagonists.
     /// </summary>
     [DataField]
-    public List<(EntityUid, string)> SelectedMinds = new();
+    public List<(EntityUid, string)> AssignedMinds = new();
 
     /// <summary>
     /// When the antag selection will occur.
@@ -37,11 +43,16 @@ public sealed partial class AntagSelectionComponent : Component
     [DataField]
     public AntagSelectionTime SelectionTime = AntagSelectionTime.PostPlayerSpawn;
 
+    /// <summary>
+    /// Cached sessions of players who are chosen yet not given the role yet.
+    /// </summary>
+    public 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.
     /// Is not serialized.
     /// </summary>
-    public HashSet<ICommonSession> SelectedSessions = new();
+    public HashSet<ICommonSession> AssignedSessions = new();
 
     /// <summary>
     /// Locale id for the name of the antag.

Content.Server/GameTicking/Rules/GameRuleSystem.Utility.cs

@@ -19,6 +19,11 @@ public abstract partial class GameRuleSystem<T> where T: IComponent
         return EntityQueryEnumerator<ActiveGameRuleComponent, T, GameRuleComponent>();
     }
 
+    protected EntityQueryEnumerator<DelayedStartRuleComponent, T, GameRuleComponent> QueryDelayedRules()
+    {
+        return EntityQueryEnumerator<DelayedStartRuleComponent, T, GameRuleComponent>();
+    }
+
     /// <summary>
     /// Queries all gamerules, regardless of if they're active or not.
     /// </summary>

Content.Server/Station/Systems/StationJobsSystem.Roundstart.cs

@@ -1,10 +1,12 @@
 using System.Linq;
 using Content.Server.Administration.Managers;
+using Content.Server.Antag;
 using Content.Server.Players.PlayTimeTracking;
 using Content.Server.Station.Components;
 using Content.Server.Station.Events;
 using Content.Shared.Preferences;
 using Content.Shared.Roles;
+using Robust.Server.Player;
 using Robust.Shared.Network;
 using Robust.Shared.Prototypes;
 using Robust.Shared.Random;
@@ -17,6 +19,8 @@ public sealed partial class StationJobsSystem
 {
     [Dependency] private readonly IPrototypeManager _prototypeManager = default!;
     [Dependency] private readonly IBanManager _banManager = default!;
+    [Dependency] private readonly IPlayerManager _playerManager = default!;
+    [Dependency] private readonly AntagSelectionSystem _antag = default!;
 
     private Dictionary<int, HashSet<string>> _jobsByWeight = default!;
     private List<int> _orderedWeights = default!;
@@ -345,6 +349,7 @@ public sealed partial class StationJobsSystem
         foreach (var (player, profile) in profiles)
         {
             var roleBans = _banManager.GetJobBans(player);
+            var antagBlocked = _antag.GetPreSelectedAntagSessions();
             var profileJobs = profile.JobPriorities.Keys.Select(k => new ProtoId<JobPrototype>(k)).ToList();
             var ev = new StationJobsGetCandidatesEvent(player, profileJobs);
             RaiseLocalEvent(ref ev);
@@ -361,6 +366,9 @@ public sealed partial class StationJobsSystem
                 if (!_prototypeManager.TryIndex(jobId, out var job))
                     continue;
 
+                if (!job.CanBeAntag && (!_playerManager.TryGetSessionById(player, out var session) || antagBlocked.Contains(session)))
+                    continue;
+
                 if (weight is not null && job.Weight != weight.Value)
                     continue;
 

Content.Shared/Antag/AntagAcceptability.cs

@@ -17,7 +17,7 @@ public enum AntagAcceptability
     /// <summary>
     /// Choose anyone
     /// </summary>
-    All
+    All,
 }
 
 public enum AntagSelectionTime : byte
@@ -28,8 +28,14 @@ public enum AntagSelectionTime : byte
     /// </summary>
     PrePlayerSpawn,
 
+    /// <summary>
+    /// Antag roles are selected to the player session before job assignment and spawning.
+    /// Unlike PrePlayerSpawn, this does not remove you from the job spawn pool.
+    /// </summary>
+    IntraPlayerSpawn,
+
     /// <summary>
     /// Antag roles get assigned after players have been assigned jobs and have spawned in.
     /// </summary>
-    PostPlayerSpawn
+    PostPlayerSpawn,
 }