.NET: Fix and re-enable flaky InputWaiter timeout test

dotnet/src/Microsoft.Agents.AI.Workflows/Execution/InputWaiter.cs

@@ -33,10 +33,24 @@ public void SignalInput()
         }
     }
 
-    public Task WaitForInputAsync(CancellationToken cancellationToken = default) => this.WaitForInputAsync(null, cancellationToken);
+    /// <summary>
+    /// Waits until input is signaled. This wait never expires; it completes only when
+    /// <see cref="SignalInput"/> is called or <paramref name="cancellationToken"/> is cancelled.
+    /// </summary>
+    /// <param name="cancellationToken">A token to cancel the wait.</param>
+    public Task WaitForInputAsync(CancellationToken cancellationToken = default) => this._inputSignal.WaitAsync(cancellationToken);
 
-    public async Task WaitForInputAsync(TimeSpan? timeout = null, CancellationToken cancellationToken = default)
+    /// <summary>
+    /// Waits until input is signaled or <paramref name="timeout"/> expires.
+    /// </summary>
+    /// <param name="timeout">The maximum time to wait for input.</param>
+    /// <param name="cancellationToken">A token to cancel the wait.</param>
+    /// <returns>
+    /// <see langword="true"/> if the wait was released by <see cref="SignalInput"/>;
+    /// <see langword="false"/> if <paramref name="timeout"/> expired first.
+    /// </returns>
+    public async Task<bool> WaitForInputAsync(TimeSpan timeout, CancellationToken cancellationToken = default)
     {
-        await this._inputSignal.WaitAsync(timeout ?? TimeSpan.FromMilliseconds(-1), cancellationToken).ConfigureAwait(false);
+        return await this._inputSignal.WaitAsync(timeout, cancellationToken).ConfigureAwait(false);
     }
 }

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/InputWaiterTests.cs

@@ -10,6 +10,12 @@ namespace Microsoft.Agents.AI.Workflows.UnitTests;
 
 public sealed class InputWaiterTests : IDisposable
 {
+    /// <summary>
+    /// Liveness backstop for waits that are expected to complete. Never the thing under test:
+    /// it is set far above the timeouts being exercised so a regression fails instead of hanging.
+    /// </summary>
+    private static readonly TimeSpan s_guardTimeout = TimeSpan.FromSeconds(30);
+
     private readonly InputWaiter _waiter = new();
 
     public void Dispose()
@@ -21,39 +27,37 @@ public void Dispose()
     [Fact]
     public async Task InputWaiter_WaitForInputAsync_CompletesAfterSignalAsync()
     {
+        // Arrange
         this._waiter.SignalInput();
 
-        // WaitForInputAsync should complete immediately since input was already signaled
-        Task waitTask = this._waiter.WaitForInputAsync(CancellationToken.None);
-        Task completed = await Task.WhenAny(waitTask, Task.Delay(TimeSpan.FromSeconds(1)));
+        // Act
+        bool signaled = await this._waiter.WaitForInputAsync(s_guardTimeout);
 
-        completed.Should().BeSameAs(waitTask, "the wait task should complete before the timeout");
-        await waitTask;
+        // Assert
+        signaled.Should().BeTrue("the already-signaled input should release the wait");
     }
 
     [Fact]
     public async Task InputWaiter_WaitForInputAsync_BlocksUntilSignaledAsync()
     {
-        // Use the no-timeout overload so that the wait can only be released by SignalInput.
-        // A finite timeout would make this test's logic racy: the component correctly
-        // honors the timeout, but if the test thread is starved of CPU time (CI load,
-        // GC pause) long enough for the timeout to fire, waitTask completes before
-        // SignalInput is called and the "should not complete before signaled" assertion
-        // flakes. Timeout behavior is covered separately below.
-        Task waitTask = this._waiter.WaitForInputAsync(CancellationToken.None);
+        // Arrange - the no-timeout overload is used so that only SignalInput can release the wait.
+        using CancellationTokenSource guard = new();
+        Task waitTask = this._waiter.WaitForInputAsync(guard.Token);
 
+        // Assert - the waiter stays blocked while no input has been signaled.
         Task completedBeforeSignal = await Task.WhenAny(waitTask, Task.Delay(100));
         completedBeforeSignal.Should().NotBeSameAs(
             waitTask,
             "the waiter should not complete before input is signaled");
 
+        // Act
         this._waiter.SignalInput();
 
-        Task completedAfterSignal = await Task.WhenAny(waitTask, Task.Delay(TimeSpan.FromSeconds(1)));
-        completedAfterSignal.Should().BeSameAs(
-            waitTask,
-            "the wait task should complete after being signaled");
+        // Armed only once the signal has released the wait, since cancelling a still-pending
+        // waiter would fail the test rather than guard it.
+        guard.CancelAfter(s_guardTimeout);
 
+        // Assert - completion alone proves the signal released the wait.
         await waitTask;
     }
 
@@ -98,33 +102,34 @@ public async Task InputWaiter_WaitForInputAsync_DoesNotCompleteWhenNotSignaledAs
     [Fact]
     public async Task InputWaiter_WaitForInputAsync_CanBeSignaledMultipleTimesSequentiallyAsync()
     {
-        // First signal/wait cycle
+        // Arrange / Act - first signal/wait cycle
         this._waiter.SignalInput();
-        await this._waiter.WaitForInputAsync(TimeSpan.FromSeconds(1));
+        bool firstSignaled = await this._waiter.WaitForInputAsync(s_guardTimeout);
 
-        // Second signal/wait cycle
+        // Arrange / Act - second signal/wait cycle
         this._waiter.SignalInput();
-        await this._waiter.WaitForInputAsync(TimeSpan.FromSeconds(1));
+        bool secondSignaled = await this._waiter.WaitForInputAsync(s_guardTimeout);
+
+        // Assert each cycle was released by its signal rather than by an expiring timeout.
+        firstSignaled.Should().BeTrue("the first signal should release the first wait");
+        secondSignaled.Should().BeTrue("the second signal should release the second wait");
     }
 
-    [Fact(Skip = "Flaky in merge_group; see https://github.com/microsoft/agent-framework/issues/7360")]
+    [Fact]
     public async Task InputWaiter_WaitForInputAsync_CompletesWhenTimeoutExpiresAsync()
     {
-        // Verify that a finite timeout releases the block even without a signal.
-        // We only assert that it *does* complete (within a generous outer bound);
-        // we intentionally do not assert that it stays blocked until the timeout,
-        // because that would re-introduce the same wall-clock flakiness
-        // described in BlocksUntilSignaledAsync (see comment on that test).
-        //
-        // The generous outer bound is itself not generous enough: on the loaded
-        // net472/windows-latest leg, thread pool starvation can delay the 300ms
-        // continuation past the 5s guard, so Task.Delay wins the race and the
-        // identity assertion below fails. Quarantined until the assertion is
-        // rewritten to not depend on scheduling latency.
-        Task waitTask = this._waiter.WaitForInputAsync(TimeSpan.FromMilliseconds(300));
-
-        Task completed = await Task.WhenAny(waitTask, Task.Delay(TimeSpan.FromSeconds(5)));
-        completed.Should().BeSameAs(waitTask, "the wait task should complete once the timeout expires");
-        await waitTask;
+        // Arrange - nothing signals this waiter, so an expiring timeout is the only thing that
+        // can release the wait, and the returned flag proves which one did. The guard only
+        // bounds a wait that never returns; it is not part of the assertion.
+        using CancellationTokenSource guard = new();
+
+        // Act
+        Task<bool> waitTask = this._waiter.WaitForInputAsync(TimeSpan.FromMilliseconds(300), guard.Token);
+        guard.CancelAfter(s_guardTimeout);
+
+        bool signaled = await waitTask;
+
+        // Assert
+        signaled.Should().BeFalse("the wait should be released by the expiring timeout rather than by a signal");
     }
 }