Mind tweaks & fixes (#21203)

Content.Shared/Mind/SharedMindSystem.cs

@@ -1,4 +1,5 @@
 using System.Diagnostics.CodeAnalysis;
+using System.Linq;
 using Content.Shared.Administration.Logs;
 using Content.Shared.Database;
 using Content.Shared.Examine;
@@ -25,7 +26,7 @@ public abstract class SharedMindSystem : EntitySystem
     [Dependency] private readonly SharedPlayerSystem _player = default!;
     [Dependency] private readonly MetaDataSystem _metadata = default!;
 
-    // This is dictionary is required to track the minds of disconnected players that may have had their entity deleted.
+    [ViewVariables]
     protected readonly Dictionary<NetUserId, EntityUid> UserMinds = new();
 
     public override void Initialize()
@@ -36,6 +37,7 @@ public abstract class SharedMindSystem : EntitySystem
         SubscribeLocalEvent<MindContainerComponent, SuicideEvent>(OnSuicide);
         SubscribeLocalEvent<VisitingMindComponent, EntityTerminatingEvent>(OnVisitingTerminating);
         SubscribeLocalEvent<RoundRestartCleanupEvent>(OnReset);
+        SubscribeLocalEvent<MindComponent, ComponentStartup>(OnMindStartup);
     }
 
     public override void Shutdown()
@@ -44,6 +46,29 @@ public abstract class SharedMindSystem : EntitySystem
         WipeAllMinds();
     }
 
+    private void OnMindStartup(EntityUid uid, MindComponent component, ComponentStartup args)
+    {
+        if (component.UserId == null)
+            return;
+
+        if (UserMinds.TryAdd(component.UserId.Value, uid))
+            return;
+
+        var existing = UserMinds[component.UserId.Value];
+        if (existing == uid)
+            return;
+
+        if (!Exists(existing))
+        {
+            Log.Error($"Found deleted entity in mind dictionary while initializing mind {ToPrettyString(uid)}");
+            UserMinds[component.UserId.Value] = uid;
+            return;
+        }
+
+        Log.Error($"Encountered a user {component.UserId} that is already assigned to a mind while initializing mind {ToPrettyString(uid)}. Ignoring user field.");
+        component.UserId = null;
+    }
+
     private void OnReset(RoundRestartCleanupEvent ev)
     {
         WipeAllMinds();
@@ -51,12 +76,22 @@ public abstract class SharedMindSystem : EntitySystem
 
     public virtual void WipeAllMinds()
     {
-        foreach (var mind in UserMinds.Values)
+        Log.Info($"Wiping all minds");
+        foreach (var mind in UserMinds.Values.ToArray())
         {
             WipeMind(mind);
         }
 
-        DebugTools.Assert(UserMinds.Count == 0);
+        if (UserMinds.Count == 0)
+            return;
+
+        foreach (var mind in UserMinds.Values)
+        {
+            if (Exists(mind))
+                Log.Error($"Failed to wipe mind: {ToPrettyString(mind)}");
+        }
+
+        UserMinds.Clear();
     }
 
     public EntityUid? GetMind(NetUserId user)
@@ -80,6 +115,26 @@ public abstract class SharedMindSystem : EntitySystem
         return false;
     }
 
+    public bool TryGetMind(NetUserId user, [NotNullWhen(true)] out Entity<MindComponent>? mind)
+    {
+        if (!TryGetMind(user, out var mindId, out var mindComp))
+        {
+            mind = null;
+            return false;
+        }
+
+        mind = (mindId.Value, mindComp);
+        return true;
+    }
+
+    public Entity<MindComponent> GetOrCreateMind(NetUserId user)
+    {
+        if (!TryGetMind(user, out var mind))
+            mind = CreateMind(user);
+
+        return mind.Value;
+    }
+
     private void OnVisitingTerminating(EntityUid uid, VisitingMindComponent component, ref EntityTerminatingEvent args)
     {
         if (component.MindId != null)
@@ -128,7 +183,7 @@ public abstract class SharedMindSystem : EntitySystem
         return null;
     }
 
-    public EntityUid CreateMind(NetUserId? userId, string? name = null)
+    public Entity<MindComponent> CreateMind(NetUserId? userId, string? name = null)
     {
         var mindId = Spawn(null, MapCoordinates.Nullspace);
         _metadata.SetEntityName(mindId, name == null ? "mind" : $"mind ({name})");
@@ -136,7 +191,7 @@ public abstract class SharedMindSystem : EntitySystem
         mind.CharacterName = name;
         SetUserId(mindId, userId, mind);
 
-        return mindId;
+        return (mindId, mind);
     }
 
     /// <summary>
@@ -195,7 +250,7 @@ public abstract class SharedMindSystem : EntitySystem
     /// Cleans up the VisitingEntity.
     /// </summary>
     /// <param name="mind"></param>
-    protected void RemoveVisitingEntity(MindComponent mind)
+    protected void RemoveVisitingEntity(EntityUid mindId, MindComponent mind)
     {
         if (mind.VisitingEntity == null)
             return;
@@ -210,6 +265,7 @@ public abstract class SharedMindSystem : EntitySystem
             RemCompDeferred(oldVisitingEnt, visitComp);
         }
 
+        Dirty(mindId, mind);
         RaiseLocalEvent(oldVisitingEnt, new MindUnvisitedMessage(), true);
     }
 
@@ -228,7 +284,7 @@ public abstract class SharedMindSystem : EntitySystem
         if (mindId == null || !Resolve(mindId.Value, ref mind, false))
             return;
 
-        TransferTo(mindId.Value, null, mind: mind);
+        TransferTo(mindId.Value, null, createGhost:false, mind: mind);
         SetUserId(mindId.Value, null, mind: mind);
     }
 
@@ -391,21 +447,6 @@ public abstract class SharedMindSystem : EntitySystem
         return TryComp(mindContainer.Mind, out role);
     }
 
-    /// <summary>
-    /// Sets the Mind's OwnedComponent and OwnedEntity
-    /// </summary>
-    /// <param name="mind">Mind to set OwnedComponent and OwnedEntity on</param>
-    /// <param name="uid">Entity owned by <paramref name="mind"/></param>
-    /// <param name="mindContainerComponent">MindContainerComponent owned by <paramref name="mind"/></param>
-    protected void SetOwnedEntity(MindComponent mind, EntityUid? uid, MindContainerComponent? mindContainerComponent)
-    {
-        if (uid != null)
-            Resolve(uid.Value, ref mindContainerComponent);
-
-        mind.OwnedEntity = uid;
-        mind.OwnedComponent = mindContainerComponent;
-    }
-
     /// <summary>
     /// Sets the Mind's UserId, Session, and updates the player's PlayerData. This should have no direct effect on the
     /// entity that any mind is connected to, except as a side effect of the fact that it may change a player's