fix(ribbon): register the engine-toggle prime marker before the prime starts

TaskMaster.Test/Ribbon/EngineToggleStateCoordinatorTests.PrimeRegistration.cs

@@ -0,0 +1,175 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Moq;
+
+namespace TaskMaster.Test.Ribbon
+{
+    /// <summary>
+    /// Regression tests for issue #944: the prime marker must be registered before the prime
+    /// starts, so a prime that completes on any thread always finds its own marker to remove and
+    /// a finished failed or canceled prime never blocks a later re-prime. A fourth partial of the
+    /// coordinator fixture, so the private <c>Harness</c> and <c>LoggedError</c> types and the
+    /// fixture constants are reused without adding any harness member.
+    /// </summary>
+    public partial class EngineToggleStateCoordinatorTests
+    {
+        #region Issue #944 — prime marker registration precedes the prime start
+
+        /// <summary>
+        /// Regression for issue #944 and the test that carries the fail-before obligation.
+        /// Invariant: the prime handle is registered before the activation read runs.
+        /// The read's setup callback runs synchronously inside the prime start, on the test
+        /// thread, and only records the handle it observes; every assertion runs after the
+        /// callback has returned, because an assertion thrown inside it would become a prime
+        /// fault. Without the fix no handle is registered during the read, so the recorded
+        /// handle is the already completed <see cref="Task.CompletedTask"/>, and the outcome is
+        /// decided by program order alone.
+        /// </summary>
+        [TestMethod]
+        public async Task GetPressed_WhenPrimeStarts_RegistersPrimeHandleBeforeActivationReadRuns()
+        {
+            // Arrange
+            var harness = new Harness();
+            var failure = new InvalidOperationException("configuration load failed");
+            Task handleSeenDuringRead = null;
+            // Initialized to true so that a callback that never runs fails the first assertion.
+            var handleCompletedDuringRead = true;
+            harness
+                .Engines.Setup(x => x.EngineActiveAsync(SpamEngine))
+                .Returns(() =>
+                {
+                    handleSeenDuringRead = harness.Coordinator.GetPrimeTask(SpamEngine);
+                    handleCompletedDuringRead = handleSeenDuringRead.IsCompleted;
+                    return Task.FromException<bool>(failure);
+                });
+
+            // Act
+            harness.Coordinator.GetPressed(SpamEngine);
+
+            // Assert
+            handleCompletedDuringRead
+                .Should()
+                .BeFalse(
+                    "the prime handle must be registered before the activation read runs, "
+                        + "so a prime that completes on any thread finds its own marker"
+                );
+            await handleSeenDuringRead;
+            harness.Errors.Should().ContainSingle("a faulted prime is reported exactly once");
+            harness
+                .Errors[0]
+                .Exception.Should()
+                .BeSameAs(failure, "the sink receives the injected exception unchanged");
+            var handleAfterward = harness.Coordinator.GetPrimeTask(SpamEngine);
+            handleAfterward
+                .Should()
+                .NotBeSameAs(
+                    handleSeenDuringRead,
+                    "a failed prime removes its marker before its handle completes"
+                );
+            handleAfterward
+                .IsCompleted.Should()
+                .BeTrue("with no marker registered the returned handle is already complete");
+        }
+
+        /// <summary>
+        /// Regression guard for issue #944: after a prime whose activation read returns an
+        /// already faulted task, a later read starts a new prime. With the fix the outcome is
+        /// deterministic; without it this test fails only when a thread-pool thread removes the
+        /// marker before it is stored, so it guards the user-visible outcome and does not carry
+        /// the fail-before obligation.
+        /// </summary>
+        [TestMethod]
+        public async Task GetPressed_AfterPrimeFaultsSynchronously_LaterReadStartsNewPrime()
+        {
+            // Arrange
+            var harness = new Harness();
+            var failure = new InvalidOperationException("configuration load failed");
+            harness
+                .Engines.SetupSequence(x => x.EngineActiveAsync(SpamEngine))
+                .Returns(Task.FromException<bool>(failure))
+                .Returns(Task.FromResult(true));
+
+            // Act
+            harness.Coordinator.GetPressed(SpamEngine);
+            await harness.Coordinator.GetPrimeTask(SpamEngine);
+            harness.Coordinator.GetPressed(SpamEngine);
+            var secondPrime = harness.Coordinator.GetPrimeTask(SpamEngine);
+            await secondPrime;
+
+            // Assert
+            harness.Engines.Verify(
+                x => x.EngineActiveAsync(SpamEngine),
+                Times.Exactly(2),
+                "a failed prime leaves no marker behind, so the later read starts a new prime"
+            );
+            harness
+                .Coordinator.GetPressed(SpamEngine)
+                .Should()
+                .BeTrue("the new prime read the engine as active and cached that value");
+            harness
+                .Invalidations.Should()
+                .Equal(
+                    new[] { SpamToggleControlId },
+                    "only the successful prime changed state to display"
+                );
+            harness.Errors.Should().ContainSingle("only the first prime failed");
+            harness
+                .Errors[0]
+                .Exception.Should()
+                .BeSameAs(failure, "the sink receives the injected exception unchanged");
+        }
+
+        /// <summary>
+        /// Regression guard for issue #944, canceled variant: after a prime whose activation read
+        /// returns an already canceled task, a later read starts a new prime. As with the faulted
+        /// variant, only the fixed code makes this outcome deterministic, so this test does not
+        /// carry the fail-before obligation.
+        /// </summary>
+        [TestMethod]
+        public async Task GetPressed_AfterPrimeIsCanceledSynchronously_LaterReadStartsNewPrime()
+        {
+            // Arrange
+            var harness = new Harness();
+            harness
+                .Engines.SetupSequence(x => x.EngineActiveAsync(SpamEngine))
+                .Returns(Task.FromCanceled<bool>(new CancellationToken(true)))
+                .Returns(Task.FromResult(true));
+
+            // Act
+            harness.Coordinator.GetPressed(SpamEngine);
+            await harness.Coordinator.GetPrimeTask(SpamEngine);
+            harness.Coordinator.GetPressed(SpamEngine);
+            var secondPrime = harness.Coordinator.GetPrimeTask(SpamEngine);
+            await secondPrime;
+
+            // Assert
+            harness.Engines.Verify(
+                x => x.EngineActiveAsync(SpamEngine),
+                Times.Exactly(2),
+                "a canceled prime leaves no marker behind, so the later read starts a new prime"
+            );
+            harness
+                .Coordinator.GetPressed(SpamEngine)
+                .Should()
+                .BeTrue("the new prime read the engine as active and cached that value");
+            harness
+                .Invalidations.Should()
+                .Equal(
+                    new[] { SpamToggleControlId },
+                    "only the successful prime changed state to display"
+                );
+            harness.Errors.Should().ContainSingle("only the first prime was canceled");
+            harness
+                .Errors[0]
+                .Exception.Should()
+                .BeAssignableTo<OperationCanceledException>(
+                    "a canceled task carries no exception to unwrap, so one is synthesized"
+                );
+        }
+
+        #endregion Issue #944 — prime marker registration precedes the prime start
+    }
+}

TaskMaster.Test/TaskMaster.Test.csproj

@@ -358,6 +358,7 @@
     <Compile Include="Ribbon\SpamManagerResetGateTests.cs" />
     <Compile Include="Ribbon\EngineToggleStateCoordinatorTests.Race.cs" />
     <Compile Include="Ribbon\EngineToggleStateCoordinatorTests.PrimeFaultOrdering.cs" />
+    <Compile Include="Ribbon\EngineToggleStateCoordinatorTests.PrimeRegistration.cs" />
     <Compile Include="Ribbon\EngineTogglePressedStateCacheTests.cs" />
     <Compile Include="Properties\AssemblyInfo.cs" />
   </ItemGroup>

TaskMaster/Ribbon/EngineToggleStateCoordinator.cs

@@ -56,8 +56,8 @@ internal sealed class EngineToggleStateCoordinator
         private readonly Action<string, Exception> _logError;
 
         /// <summary>
-        /// Serializes the at-most-one-prime decision. Held only across a dictionary probe and a
-        /// task start; no await occurs inside it.
+        /// Serializes the at-most-one-prime decision. Held only across a dictionary probe, the
+        /// marker registration, and the start of the prime; no await occurs inside it.
         /// </summary>
         private readonly object _primeGate = new object();
 
@@ -70,9 +70,10 @@ internal sealed class EngineToggleStateCoordinator
             new EngineTogglePressedStateCache();
 
         /// <summary>
-        /// The in-flight — or most recently completed — prime per engine key. Its presence is the
-        /// at-most-one-prime guard; its value is the test-observable handle returned by
-        /// <see cref="GetPrimeTask"/>.
+        /// The registration marker per engine key: registered before the prime starts, removed by
+        /// <see cref="CompletePrime"/> when the prime faults or is canceled, and retained after a
+        /// successful prime. Its presence is the at-most-one-prime guard; its value is the
+        /// test-observable handle returned by <see cref="GetPrimeTask"/>.
         /// </summary>
         private readonly ConcurrentDictionary<string, Task> _primeTasks = new ConcurrentDictionary<
             string,
@@ -275,7 +276,15 @@ private void StartPrimeIfNeeded(string engineName, string controlId)
                     return;
                 }
 
-                _primeTasks[engineName] = StartObservedPrime(engines, engineName, controlId);
+                // Registration precedes the start (issue #944): a prime can complete on any
+                // thread, including before StartObservedPrime returns, and it must always find
+                // its own marker to remove; registering afterwards let a finished prime's
+                // removal run first and leave a stale marker that blocked every later re-prime.
+                var marker = new TaskCompletionSource<bool>(
+                    TaskCreationOptions.RunContinuationsAsynchronously
+                );
+                _primeTasks[engineName] = marker.Task;
+                StartObservedPrime(engines, engineName, controlId, marker);
             }
         }
 
@@ -286,18 +295,31 @@ private void StartPrimeIfNeeded(string engineName, string controlId)
         /// The observer is a continuation rather than a <c>catch</c> clause, so this type keeps
         /// exactly one <c>catch</c> — the click boundary. Reading
         /// <see cref="Task.Exception"/> inside <see cref="CompletePrime"/> marks the fault
-        /// observed, so no unobserved task remains. The returned continuation task always
-        /// completes successfully, which is what makes it safe for a test to await.
+        /// observed, so no unobserved task remains. The continuation task itself is discarded;
+        /// the value a test awaits is the marker, which the continuation completes only through
+        /// <c>SetResult</c> in a <c>finally</c> after <see cref="CompletePrime"/> exits, so it
+        /// never faults or cancels.
         /// </remarks>
-        private Task StartObservedPrime(
+        private void StartObservedPrime(
             IAppItemEngines engines,
             string engineName,
-            string controlId
+            string controlId,
+            TaskCompletionSource<bool> marker
         )
         {
-            return ApplyPrimeAsync(engines, engineName, controlId)
+            _ = ApplyPrimeAsync(engines, engineName, controlId)
                 .ContinueWith(
-                    completed => CompletePrime(completed, engineName),
+                    completed =>
+                    {
+                        try
+                        {
+                            CompletePrime(completed, engineName);
+                        }
+                        finally
+                        {
+                            marker.SetResult(true);
+                        }
+                    },
                     CancellationToken.None,
                     TaskContinuationOptions.None,
                     TaskScheduler.Default