fix: Domain reload waits complete during recovery

Assets/Tests/Editor/UnityCliLoopServerControllerStartupLockTests.cs

@@ -106,6 +106,165 @@ public async Task ScheduleStartupRecovery_WhenRecoveryIsAsync_KeepsTaskIncomplet
             Assert.That(service.RecoveryTask, Is.Null);
         }
 
+        [Test]
+        public async Task ScheduleTrackedRecovery_WhenRecoveryAlreadyRunning_ReturnsCurrentTaskWithoutStartingAnotherRecovery()
+        {
+            // Tests that duplicate tracked recovery triggers join the active recovery instead of starting a second loop.
+            int recoveryStartCount = 0;
+            TaskCompletionSource<bool> firstRecoveryCompletionSource = new();
+            UnityCliLoopServerRecoveryTrackingService service = CreateRecoveryTrackingService();
+
+            Task firstTask = service.ScheduleTrackedRecovery(() =>
+            {
+                recoveryStartCount++;
+                return firstRecoveryCompletionSource.Task;
+            });
+            Task secondTask = service.ScheduleTrackedRecovery(() =>
+            {
+                recoveryStartCount++;
+                return Task.CompletedTask;
+            });
+
+            Assert.That(secondTask, Is.SameAs(firstTask));
+            Assert.That(recoveryStartCount, Is.EqualTo(1));
+
+            firstRecoveryCompletionSource.SetResult(true);
+            await firstTask;
+        }
+
+        [Test]
+        public async Task ScheduleStartupRecovery_WhenTrackedRecoveryIsRunning_ReturnsCurrentTaskWithoutSchedulingStartupRecovery()
+        {
+            // Tests that startup recovery joins an active tracked recovery without registering a delay call.
+            int scheduledActionCount = 0;
+            TaskCompletionSource<bool> trackedRecoveryCompletionSource = new();
+            UnityCliLoopServerRecoveryTrackingService service = CreateRecoveryTrackingService();
+
+            Task trackedTask = service.ScheduleTrackedRecovery(() => trackedRecoveryCompletionSource.Task);
+            Task startupTask = service.ScheduleStartupRecovery(
+                action => scheduledActionCount++,
+                () => Task.CompletedTask);
+
+            Assert.That(startupTask, Is.SameAs(trackedTask));
+            Assert.That(scheduledActionCount, Is.EqualTo(0));
+
+            trackedRecoveryCompletionSource.SetResult(true);
+            await trackedTask;
+        }
+
+        [Test]
+        public void ScheduleStartupRecovery_WhenStartupRecoveryIsPending_ReturnsCurrentTaskWithoutSchedulingAnotherRecovery()
+        {
+            // Tests that duplicate startup recovery triggers join the active startup placeholder.
+            int scheduledActionCount = 0;
+            UnityCliLoopServerRecoveryTrackingService service = CreateRecoveryTrackingService();
+
+            Task firstTask = service.ScheduleStartupRecovery(
+                action => scheduledActionCount++,
+                () => Task.CompletedTask);
+            Task secondTask = service.ScheduleStartupRecovery(
+                action => scheduledActionCount++,
+                () => Task.CompletedTask);
+
+            Assert.That(secondTask, Is.SameAs(firstTask));
+            Assert.That(scheduledActionCount, Is.EqualTo(1));
+        }
+
+        [Test]
+        public async Task ScheduleTrackedRecovery_WhenStartupRecoveryIsPending_ReplacesStartupPlaceholderWithTrackedRecovery()
+        {
+            // Tests that after-domain-reload recovery is not dropped behind startup recovery.
+            System.Action scheduledAction = null;
+            int trackedRecoveryStartCount = 0;
+            TaskCompletionSource<bool> startupRecoveryCompletionSource = new();
+            TaskCompletionSource<bool> trackedRecoveryCompletionSource = new();
+            UnityCliLoopServerRecoveryTrackingService service = CreateRecoveryTrackingService();
+
+            Task startupTask = service.ScheduleStartupRecovery(
+                action => scheduledAction = action,
+                () => startupRecoveryCompletionSource.Task);
+            Task trackedTask = service.ScheduleTrackedRecovery(() =>
+            {
+                trackedRecoveryStartCount++;
+                return trackedRecoveryCompletionSource.Task;
+            });
+
+            Assert.That(trackedTask, Is.Not.SameAs(startupTask));
+            Assert.That(trackedTask, Is.SameAs(service.RecoveryTask));
+            Assert.That(trackedRecoveryStartCount, Is.EqualTo(1));
+
+            scheduledAction();
+            startupRecoveryCompletionSource.SetResult(true);
+            await startupTask;
+            Assert.That(service.RecoveryTask, Is.SameAs(trackedTask));
+
+            trackedRecoveryCompletionSource.SetResult(true);
+            await trackedTask;
+            Assert.That(service.RecoveryTask, Is.Null);
+        }
+
+        [Test]
+        public async Task ScheduleTrackedRecovery_WhenPreviousRecoveryCompleted_StartsNewRecovery()
+        {
+            // Tests that completed recoveries do not block a later recovery trigger.
+            int recoveryStartCount = 0;
+            TaskCompletionSource<bool> firstCompletionSource = new();
+            TaskCompletionSource<bool> secondCompletionSource = new();
+            UnityCliLoopServerRecoveryTrackingService service = CreateRecoveryTrackingService();
+
+            Task firstTask = service.ScheduleTrackedRecovery(() =>
+            {
+                recoveryStartCount++;
+                return firstCompletionSource.Task;
+            });
+            firstCompletionSource.SetResult(true);
+            await firstTask;
+            Task secondTask = service.ScheduleTrackedRecovery(() =>
+            {
+                recoveryStartCount++;
+                return secondCompletionSource.Task;
+            });
+            secondCompletionSource.SetResult(true);
+            await secondTask;
+
+            Assert.That(secondTask, Is.Not.SameAs(firstTask));
+            Assert.That(recoveryStartCount, Is.EqualTo(2));
+        }
+
+        [Test]
+        public async Task ScheduleTrackedRecovery_WhenPreviousRecoveryFaulted_StartsNewRecovery()
+        {
+            // Tests that a faulted recovery does not block a later recovery trigger.
+            UnityEngine.TestTools.LogAssert.Expect(
+                UnityEngine.LogType.Error,
+                "[UnityCliLoop] Unity CLI Loop server recovery failed before the bridge became ready. first failure");
+            int firstRecoveryAttempts = 0;
+            bool secondRecoveryStarted = false;
+            UnityCliLoopServerRecoveryTrackingService service = CreateRecoveryTrackingService(
+                (delayMilliseconds, ct) => Task.CompletedTask);
+
+            Task firstTask = service.ScheduleTrackedRecovery(() =>
+            {
+                firstRecoveryAttempts++;
+                throw new System.TimeoutException("first failure");
+            });
+            Assert.ThrowsAsync<System.InvalidOperationException>(async () => await firstTask);
+            Task secondTask = service.ScheduleTrackedRecovery(() =>
+            {
+                secondRecoveryStarted = true;
+                TaskCompletionSource<bool> completionSource = new();
+                completionSource.SetResult(true);
+                return completionSource.Task;
+            });
+            await secondTask;
+
+            Assert.That(secondTask, Is.Not.SameAs(firstTask));
+            Assert.That(
+                firstRecoveryAttempts,
+                Is.EqualTo(UnityCliLoopServerConfig.RECOVERY_RETRY_DELAYS_MS.Length + 1));
+            Assert.That(secondRecoveryStarted, Is.True);
+        }
+
         [Test]
         public async Task StartRecoveryIfNeededAsync_WhenReadinessSucceeds_ShouldPublishServerStartedEvent()
         {

Packages/src/Editor/Infrastructure/Server/UnityCliLoopServerRecoveryTrackingService.cs

@@ -16,6 +16,7 @@ internal sealed class UnityCliLoopServerRecoveryTrackingService
         private readonly ISessionFlagsRepository _sessionFlagsRepository;
         private readonly Func<int, CancellationToken, Task> _waitBeforeRecoveryRetryAsync;
         private Task _currentRecoveryTask;
+        private bool _isCurrentRecoveryTracked;
 
         internal UnityCliLoopServerRecoveryTrackingService(
             ISessionFlagsRepository sessionFlagsRepository,
@@ -39,8 +40,16 @@ internal Task ScheduleStartupRecovery(
             Debug.Assert(scheduleDelayCall != null, "scheduleDelayCall must not be null");
             Debug.Assert(restoreServerState != null, "restoreServerState must not be null");
 
-            TaskCompletionSource<bool> scheduledRecoveryCompletionSource = new();
-            _currentRecoveryTask = scheduledRecoveryCompletionSource.Task;
+            TaskCompletionSource<bool> scheduledRecoveryCompletionSource = null;
+            Task scheduledRecoveryTask = ScheduleRecoveryTask(() =>
+            {
+                scheduledRecoveryCompletionSource = new TaskCompletionSource<bool>();
+                return scheduledRecoveryCompletionSource.Task;
+            }, isTrackedRecovery: false);
+            if (scheduledRecoveryCompletionSource == null)
+            {
+                return scheduledRecoveryTask;
+            }
 
             scheduleDelayCall(() =>
             {
@@ -67,7 +76,7 @@ internal Task ScheduleStartupRecovery(
                 }, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.FromCurrentSynchronizationContext());
             });
 
-            return scheduledRecoveryCompletionSource.Task;
+            return scheduledRecoveryTask;
         }
 
         private void CompleteScheduledStartupRecovery(
@@ -77,6 +86,7 @@ private void CompleteScheduledStartupRecovery(
             if (ReferenceEquals(_currentRecoveryTask, scheduledRecoveryCompletionSource.Task))
             {
                 _currentRecoveryTask = null;
+                _isCurrentRecoveryTracked = false;
             }
 
             if (restoreTask.IsCanceled)
@@ -104,12 +114,63 @@ internal Task ScheduleTrackedRecovery(Func<Task> recoveryAction)
         {
             Debug.Assert(recoveryAction != null, "recoveryAction must not be null");
 
-            Task recoveryTask = ExecuteTrackedRecoveryAsync(recoveryAction);
+            Task trackedRecoveryTask = null;
+            Task recoveryTask = ScheduleRecoveryTask(() =>
+            {
+                trackedRecoveryTask = ExecuteTrackedRecoveryAsync(recoveryAction);
+                return trackedRecoveryTask;
+            }, isTrackedRecovery: true);
+            if (trackedRecoveryTask != null)
+            {
+                _ = ClearTrackedRecoveryWhenCompleteAsync(trackedRecoveryTask);
+            }
+
+            return recoveryTask;
+        }
+
+        private Task ScheduleRecoveryTask(Func<Task> createRecoveryTask, bool isTrackedRecovery)
+        {
+            Debug.Assert(createRecoveryTask != null, "createRecoveryTask must not be null");
+
+            Task activeRecoveryTask = GetActiveRecoveryTask();
+            if (activeRecoveryTask != null && ShouldJoinActiveRecovery(isTrackedRecovery))
+            {
+                // Why: startup recovery has no domain-reload bookkeeping, and duplicate tracked
+                // recovery loops recheck current server state, so these triggers add no new work.
+                return activeRecoveryTask;
+            }
+
+            Task recoveryTask = createRecoveryTask();
             _currentRecoveryTask = recoveryTask;
-            _ = ClearTrackedRecoveryWhenCompleteAsync(recoveryTask);
+            _isCurrentRecoveryTracked = isTrackedRecovery;
             return recoveryTask;
         }
 
+        private bool ShouldJoinActiveRecovery(bool isTrackedRecovery)
+        {
+            if (!isTrackedRecovery)
+            {
+                // Why: startup restore is an idempotent IfNeeded recovery path, so any active
+                // recovery task already owns the same server-state restoration work.
+                return true;
+            }
+
+            // Why: after-domain-reload tracked recovery carries CompleteDomainReload bookkeeping,
+            // so it must supersede a startup placeholder instead of joining it.
+            return _isCurrentRecoveryTracked;
+        }
+
+        private Task GetActiveRecoveryTask()
+        {
+            Task currentRecoveryTask = _currentRecoveryTask;
+            if (currentRecoveryTask == null || currentRecoveryTask.IsCompleted)
+            {
+                return null;
+            }
+
+            return currentRecoveryTask;
+        }
+
         /// <summary>
         /// Runs a recovery action, retrying with backoff so one transient failure
         /// (e.g. readiness timeout during a heavy import) does not leave the server
@@ -176,6 +237,7 @@ private async Task ClearTrackedRecoveryWhenCompleteAsync(Task recoveryTask)
                 if (ReferenceEquals(_currentRecoveryTask, recoveryTask))
                 {
                     _currentRecoveryTask = null;
+                    _isCurrentRecoveryTracked = false;
                 }
             }
         }