Reduce absolute time dependencies in tests (#1649)

tests/ModelContextProtocol.Tests/Server/McpServerTaskTests.cs

@@ -1,7 +1,6 @@
 using ModelContextProtocol.Protocol;
 using ModelContextProtocol.Server;
 using Microsoft.Extensions.DependencyInjection;
-using System.Collections.Concurrent;
 using System.Runtime.InteropServices;
 using System.Text.Json;
 using System.Text.Json.Nodes;
@@ -179,17 +178,25 @@ public class McpServerTaskTests : ClientServerTestBase
         await using var client = await CreateMcpClientForServer();
         var ct = TestContext.Current.CancellationToken;
 
+        var failedTask = new TaskCompletionSource<bool>();
+
+        // Run failure task once the task from the tool call is created
+        _taskStore.OnTaskCreated += taskId =>
+        {
             _ = Task.Run(async () =>
             {
                 await Task.Delay(100, ct);
-            var taskId = _taskStore.GetAllTaskIds().Single();
                 _taskStore.FailTask(taskId, JsonElement.Parse("""{"code":-32000,"message":"something went wrong"}"""));
+                failedTask.SetResult(true);
             }, ct);
+        };
 
         await Assert.ThrowsAsync<McpException>(async () =>
             await client.CallToolAsync(
                 new CallToolRequestParams { Name = "async-tool" },
                 ct));
+
+        Assert.True(await failedTask.Task);
     }
 
     [Fact]
@@ -198,17 +205,25 @@ public class McpServerTaskTests : ClientServerTestBase
         await using var client = await CreateMcpClientForServer();
         var ct = TestContext.Current.CancellationToken;
 
-        _ = Task.Run(async () =>
+        var cancelledTask = new TaskCompletionSource<bool>();
+
+        // Run cancellation task once the task from the tool call is created
+        _taskStore.OnTaskCreated += taskId =>
+        {
+            Task.Run(async () =>
             {
                 await Task.Delay(100, ct);
-            var taskId = _taskStore.GetAllTaskIds().Single();
                 _taskStore.CancelTask(taskId);
+                cancelledTask.SetResult(true);
             }, ct);
+        };
 
         await Assert.ThrowsAsync<OperationCanceledException>(async () =>
             await client.CallToolAsync(
                 new CallToolRequestParams { Name = "async-tool" },
                 ct));
+
+        Assert.True(await cancelledTask.Task);
     }
 
     [Fact]
@@ -538,23 +553,39 @@ public class McpServerTaskTests : ClientServerTestBase
     /// </summary>
     private sealed class InMemoryTaskStore
     {
-        private readonly ConcurrentDictionary<string, TaskEntry> _tasks = new();
+        private readonly Dictionary<string, TaskEntry> _tasks = new();
+
+        internal Action<string>? OnTaskCreated;
 
         public string CreateTask(McpTaskStatus initialStatus = McpTaskStatus.Working)
         {
             var taskId = Guid.NewGuid().ToString("N");
+            lock (_tasks)
+            {
                 _tasks[taskId] = new TaskEntry
                 {
                     Status = initialStatus,
                     CreatedAt = DateTimeOffset.UtcNow,
                     LastUpdatedAt = DateTimeOffset.UtcNow,
                 };
+            }
+
+            OnTaskCreated?.Invoke(taskId);
+
             return taskId;
         }
 
-        public IEnumerable<string> GetAllTaskIds() => _tasks.Keys;
+        public IEnumerable<string> GetAllTaskIds()
+        {
+            lock (_tasks)
+            {
+                return _tasks.Keys.ToArray();
+            }
+        }
 
         public GetTaskResult GetTask(string taskId)
+        {
+            lock (_tasks)
             {
                 if (!_tasks.TryGetValue(taskId, out var entry))
                 {
@@ -600,44 +631,57 @@ public class McpServerTaskTests : ClientServerTestBase
                     _ => throw new InvalidOperationException($"Unexpected status: {entry.Status}")
                 };
             }
+        }
 
         public void CompleteTask(string taskId, CallToolResult result)
+        {
+            lock (_tasks)
             {
                 if (_tasks.TryGetValue(taskId, out var entry))
                 {
-                entry.Status = McpTaskStatus.Completed;
                     entry.Result = result;
                     entry.LastUpdatedAt = DateTimeOffset.UtcNow;
+                    entry.Status = McpTaskStatus.Completed;
+                }
             }
         }
 
         public void FailTask(string taskId, JsonElement error)
+        {
+            lock (_tasks)
             {
                 if (_tasks.TryGetValue(taskId, out var entry))
                 {
-                entry.Status = McpTaskStatus.Failed;
                     entry.Error = error;
                     entry.LastUpdatedAt = DateTimeOffset.UtcNow;
+                    entry.Status = McpTaskStatus.Failed;
+                }
             }
         }
 
         public void CancelTask(string taskId)
+        {
+            lock (_tasks)
             {
                 if (_tasks.TryGetValue(taskId, out var entry))
                 {
-                entry.Status = McpTaskStatus.Cancelled;
                     entry.LastUpdatedAt = DateTimeOffset.UtcNow;
+                    entry.Status = McpTaskStatus.Cancelled;
+                }
             }
         }
 
         public void ProvideInput(string taskId, IDictionary<string, InputResponse> inputResponses)
+        {
+            lock (_tasks)
             {
                 if (_tasks.TryGetValue(taskId, out var entry))
                 {
                     entry.InputResponses = inputResponses;
+                    entry.LastUpdatedAt = DateTimeOffset.UtcNow;
                     // Transition back to working after receiving input
                     entry.Status = McpTaskStatus.Working;
-                entry.LastUpdatedAt = DateTimeOffset.UtcNow;
+                }
             }
         }
 

tests/ModelContextProtocol.Tests/Server/TaskPollStuckDetectorTests.cs

@@ -15,6 +15,8 @@ namespace ModelContextProtocol.Tests.Server;
 /// </summary>
 public class TaskPollStuckDetectorTests : ClientServerTestBase
 {
+    private int _pollCount = 0;
+
     public TaskPollStuckDetectorTests(ITestOutputHelper testOutputHelper) : base(testOutputHelper)
     {
 #if !NET
@@ -48,6 +50,8 @@ public class TaskPollStuckDetectorTests : ClientServerTestBase
             // misbehaving-server condition the stuck-detector exists to break out of.
             options.Handlers.GetTaskHandler = (context, cancellationToken) =>
             {
+                Interlocked.Increment(ref _pollCount);
+
                 return new ValueTask<GetTaskResult>(new InputRequiredTaskResult
                 {
                     TaskId = context.Params!.TaskId,
@@ -77,19 +81,13 @@ public class TaskPollStuckDetectorTests : ClientServerTestBase
         await using var client = await CreateMcpClientForServer();
         var ct = TestContext.Current.CancellationToken;
 
-        var sw = System.Diagnostics.Stopwatch.StartNew();
-
         var ex = await Assert.ThrowsAsync<McpException>(async () =>
             await client.CallToolAsync(new CallToolRequestParams { Name = "any-tool" }, ct));
 
-        sw.Stop();
-
         Assert.Contains(McpTaskStatus.InputRequired.ToString(), ex.Message);
         Assert.Contains("consecutive polls", ex.Message);
 
-        // 60 polls × 5ms ≈ 300ms; allow generous slack for CI.
+        Assert.Equal(60, _pollCount);
-        Assert.True(sw.Elapsed < TimeSpan.FromSeconds(10),
-            $"Stuck-detector should give up promptly but took {sw.Elapsed}.");
     }
 
     [Fact]
@@ -111,6 +109,7 @@ public class TaskPollStuckDetectorTests : ClientServerTestBase
         // The message embeds the configured threshold, which is the strongest signal that the
         // option value (not the 60-default constant) is what governed the loop.
         Assert.Contains($"{CustomThreshold} consecutive polls", ex.Message);
+        Assert.Equal(CustomThreshold, _pollCount);
     }
 
     [Theory]