From 39710e7937bf9cc79a080b3755e6a4d74c5c7537 Mon Sep 17 00:00:00 2001
From: lL1l1 <82986251+lL1l1@users.noreply.github.com>
Date: Sun, 22 Sep 2024 12:54:30 -0700
Subject: [PATCH] Fix the "assist to upgrade" and "assist to unpause" settings
 conflicting and not pausing upgraded units during lag (#6446)

---
 changelog/snippets/fix.6446.md |   1 +
 lua/ui/game/commandmode.lua    | 105 ++++++++++++++++++---------------
 2 files changed, 60 insertions(+), 46 deletions(-)
 create mode 100644 changelog/snippets/fix.6446.md

diff --git a/changelog/snippets/fix.6446.md b/changelog/snippets/fix.6446.md
new file mode 100644
index 0000000000..7a6ff32a0b
--- /dev/null
+++ b/changelog/snippets/fix.6446.md
@@ -0,0 +1 @@
+- (#6446) Fix the "assist to upgrade" and "assist to unpause" game options conflicting and causing upgrading units to unpause during lag spikes.
diff --git a/lua/ui/game/commandmode.lua b/lua/ui/game/commandmode.lua
index 056f1a46e3..f893caf6ad 100644
--- a/lua/ui/game/commandmode.lua
+++ b/lua/ui/game/commandmode.lua
@@ -406,61 +406,74 @@ local function OnGuardUpgrade(guardees, unit)
     end
 end
 
+--- Thread to keep track of when to unpause, 
+--- logic is a bit convoluted but guarantees that we still have access to the user units as the game progresses
+---@param targetId EntityId
+local function UnpauseThread(targetId)
+    WaitTicks(10)
+    local target = GetUnitById(targetId)
+    while target do
+        local candidates = target.ThreadUnpauseCandidates
+        if (candidates and not table.empty(candidates)) then
+            for id, _ in candidates do
+                local engineer = GetUnitById(id)
+                -- check if it is idle instead of its guarded entity to allow queuing orders before the assist command
+                if engineer and not engineer:IsIdle() then
+                    -- ensure the target focus exists, since this thread may be targeted at something that is not building anything,
+                    -- but might start to build something after some network delay, which you won't want to unpause due to `nil == nil`
+                    local targetFocus = target:GetFocus()
+                    if targetFocus and targetFocus == engineer:GetFocus() then
+                        target.ThreadUnpauseCandidates = nil
+                        target.ThreadUnpause = nil
+                        SetPaused({ target }, false)
+                        return
+                    end
+                else
+                    -- engineer is idle, died, we switch armies, ...
+                    candidates[id] = nil
+                end
+            end
+        else
+            target.ThreadUnpauseCandidates = nil
+            target.ThreadUnpause = nil
+            return
+        end
+
+        WaitTicks(10)
+        target = GetUnitById(targetId)
+    end
+end
+
 ---@param guardees UserUnit[]
 ---@param target UserUnit
 local function OnGuardUnpause(guardees, target)
     local prefs = Prefs.GetFieldFromCurrentProfile('options').assist_to_unpause
-    if prefs == 'On' or
-        (
-        prefs == 'ExtractorsAndRadars' and
-            EntityCategoryContains((categories.MASSEXTRACTION + categories.RADAR) * categories.STRUCTURE, target))
-    then
-
-        -- start a single thread to keep track of when to unpause, logic feels a bit convoluted
-        -- but that is purely to guarantee that we still have access to the user units as the
-        -- game progresses
-        if not target.ThreadUnpause then
-            local id = target:GetEntityId()
-            target.ThreadUnpause = ForkThread(
-                function()
-                    WaitSeconds(1.0)
-                    local target = GetUnitById(id)
-                    while target do
-                        local candidates = target.ThreadUnpauseCandidates
-                        if (candidates and not table.empty(candidates)) then
-                            for id, _ in candidates do
-                                local engineer = GetUnitById(id)
-                                if engineer and not engineer:IsIdle() then
-                                    local focus = engineer:GetFocus()
-                                    if focus == target:GetFocus() then
-                                        target.ThreadUnpauseCandidates = nil
-                                        target.ThreadUnpause = nil
-                                        SetPaused({ target }, false)
-                                        break
-                                    end
-                                    -- engineer is idle, died, we switch armies, ...
-                                else
-                                    candidates[id] = nil
-                                end
-                            end
-                        else
-                            target.ThreadUnpauseCandidates = nil
-                            target.ThreadUnpause = nil
-                            break
-                        end
-
-                        WaitSeconds(1.0)
-                        target = GetUnitById(id)
-                    end
-                end
+    local bp = __blueprints[target:GetUnitId()]
+    -- only create the unpause thread for units that have the ability to unpause
+    if  (
+            prefs == 'On' and 
+            (
+                EntityCategoryContains(categories.REPAIR + categories.FACTORY + categories.SILO, target)  -- REPAIR includes mantis and harbs, compared to ENGINEER category
+                or (bp.General.UpgradesTo and bp.General.UpgradesTo ~= '') -- upgradeables can also be assisted
             )
-        end
-
-        -- add these to keep track
+        )
+        or
+        (   
+            prefs == 'ExtractorsAndRadars'
+            and EntityCategoryContains((categories.MASSEXTRACTION + categories.RADAR) * categories.STRUCTURE, target) 
+            and (bp.General.UpgradesTo and bp.General.UpgradesTo ~= '') -- use `and` to make sure the mex/radar is upgradeable
+        )
+    then
+        -- save the guardees' entity ids to keep track of in the unpause thread
         target.ThreadUnpauseCandidates = target.ThreadUnpauseCandidates or {}
         for k, guardee in guardees do
             target.ThreadUnpauseCandidates[guardee:GetEntityId()] = true
         end
+
+        -- start a single thread to keep track of when to unpause
+        if not target.ThreadUnpause then
+            target.ThreadUnpause = ForkThread(UnpauseThread, target:GetEntityId())
+        end
     end
 end