Use TimeProvider in InMemoryMcpTaskStore for deterministic testing (#1303)

src/ModelContextProtocol.Core/Server/InMemoryMcpTaskStore.cs

@@ -35,7 +35,11 @@ public sealed class InMemoryMcpTaskStore : IMcpTaskStore, IDisposable
     private readonly TimeSpan? _defaultTtl;
     private readonly TimeSpan? _maxTtl;
     private readonly TimeSpan _pollInterval;
+#if MCP_TEST_TIME_PROVIDER
+    private readonly ITimer? _cleanupTimer;
+#else
     private readonly Timer? _cleanupTimer;
+#endif
     private readonly int _pageSize;
     private readonly int? _maxTasks;
     private readonly int? _maxTasksPerSession;
@@ -134,7 +138,11 @@ public sealed class InMemoryMcpTaskStore : IMcpTaskStore, IDisposable
         cleanupInterval ??= TimeSpan.FromMinutes(1);
         if (cleanupInterval.Value != Timeout.InfiniteTimeSpan)
         {
+#if MCP_TEST_TIME_PROVIDER
+            _cleanupTimer = _timeProvider.CreateTimer(CleanupExpiredTasks, null, cleanupInterval.Value, cleanupInterval.Value);
+#else
             _cleanupTimer = new Timer(CleanupExpiredTasks, null, cleanupInterval.Value, cleanupInterval.Value);
+#endif
         }
     }
 

tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs

@@ -235,14 +235,24 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
     [Fact]
     public async Task UpdateTaskStatusAsync_UpdatesLastUpdatedAt()
     {
-        // Arrange
+        // Arrange - Use FakeTimeProvider for deterministic testing
-        using var store = new InMemoryMcpTaskStore();
+        var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
+        using var store = new TestInMemoryMcpTaskStore(
+            defaultTtl: null,
+            maxTtl: null,
+            pollInterval: null,
+            cleanupInterval: Timeout.InfiniteTimeSpan,
+            pageSize: 100,
+            maxTasks: null,
+            maxTasksPerSession: null,
+            timeProvider: fakeTime);
+            
         var metadata = new McpTaskMetadata();
         var task = await store.CreateTaskAsync(metadata, new RequestId("req-1"), new JsonRpcRequest { Method = "test" }, null, TestContext.Current.CancellationToken);
         var originalTimestamp = task.LastUpdatedAt;
 
-        // Wait a bit to ensure timestamp changes
+        // Advance time to ensure timestamp changes
-        await Task.Delay(10, TestContext.Current.CancellationToken);
+        fakeTime.Advance(TimeSpan.FromMilliseconds(10));
 
         // Act
         await store.UpdateTaskStatusAsync(task.TaskId, McpTaskStatus.Working, null, null, TestContext.Current.CancellationToken);
@@ -460,16 +470,28 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
     [Fact]
     public async Task Dispose_StopsCleanupTimer()
     {
-        // Arrange
+        // Arrange - Use FakeTimeProvider for deterministic testing
-        var store = new InMemoryMcpTaskStore(cleanupInterval: TimeSpan.FromMilliseconds(100));
+        var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
-        var metadata = new McpTaskMetadata { TimeToLive = TimeSpan.FromMilliseconds(100) }; // Very short TTL
+        var cleanupInterval = TimeSpan.FromMilliseconds(100);
+        
+        var store = new TestInMemoryMcpTaskStore(
+            defaultTtl: null,
+            maxTtl: null,
+            pollInterval: null,
+            cleanupInterval: cleanupInterval,
+            pageSize: 100,
+            maxTasks: null,
+            maxTasksPerSession: null,
+            timeProvider: fakeTime);
+            
+        var metadata = new McpTaskMetadata { TimeToLive = TimeSpan.FromMilliseconds(100) };
         await store.CreateTaskAsync(metadata, new RequestId("req-1"), new JsonRpcRequest { Method = "test" }, null, TestContext.Current.CancellationToken);
 
         // Act
         store.Dispose();
 
-        // Wait longer than cleanup interval
+        // Advance time - timer should not fire after dispose
-        await Task.Delay(300, TestContext.Current.CancellationToken);
+        fakeTime.Advance(TimeSpan.FromTicks(cleanupInterval.Ticks * 3));
 
         // Assert - Store should still be accessible after dispose (no exceptions)
         // The cleanup timer should have stopped
@@ -479,17 +501,33 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
     [Fact]
     public async Task CleanupExpiredTasks_RemovesExpiredTasks()
     {
-        // Arrange
+        // Arrange - Use FakeTimeProvider for deterministic testing
-        using var store = new InMemoryMcpTaskStore(cleanupInterval: TimeSpan.FromMilliseconds(50));
+        var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
-        var metadata = new McpTaskMetadata { TimeToLive = TimeSpan.FromMilliseconds(100) }; // 100ms TTL
+        var cleanupInterval = TimeSpan.FromMilliseconds(50);
+        var ttl = TimeSpan.FromMilliseconds(100);
+        
+        using var store = new TestInMemoryMcpTaskStore(
+            defaultTtl: null,
+            maxTtl: null,
+            pollInterval: null,
+            cleanupInterval: cleanupInterval,
+            pageSize: 100,
+            maxTasks: null,
+            maxTasksPerSession: null,
+            timeProvider: fakeTime);
+            
+        var metadata = new McpTaskMetadata { TimeToLive = ttl };
         var task = await store.CreateTaskAsync(metadata, new RequestId("req-1"), new JsonRpcRequest { Method = "test" }, null, TestContext.Current.CancellationToken);
 
         // Verify task exists initially
         var resultBefore = await store.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
         Assert.Single(resultBefore.Tasks);
 
-        // Wait for task to expire and cleanup timer to run (wait for at least 3 cleanup cycles)
+        // Advance time past the TTL to make task expired
-        await Task.Delay(250, TestContext.Current.CancellationToken);
+        fakeTime.Advance(ttl + TimeSpan.FromMilliseconds(1));
+        
+        // Trigger cleanup by advancing time past cleanup interval
+        fakeTime.Advance(cleanupInterval);
 
         // Act - List tasks to verify cleanup happened
         var resultAfter = await store.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
@@ -787,8 +825,18 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
     [Fact]
     public async Task ListTasksAsync_ConsistentWithExpiredTasksRemovedBetweenPages()
     {
-        // Arrange - Use TTL of 1 second
+        // Arrange - Use FakeTimeProvider for deterministic testing
-        using var store = new InMemoryMcpTaskStore(defaultTtl: TimeSpan.FromSeconds(1), pageSize: 5, cleanupInterval: Timeout.InfiniteTimeSpan);
+        var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
+        var ttl = TimeSpan.FromSeconds(1);
+        using var store = new TestInMemoryMcpTaskStore(
+            defaultTtl: ttl,
+            maxTtl: null,
+            pollInterval: null,
+            cleanupInterval: Timeout.InfiniteTimeSpan,
+            pageSize: 5,
+            maxTasks: null,
+            maxTasksPerSession: null,
+            timeProvider: fakeTime);
 
         // Create 15 tasks
         for (int i = 0; i < 15; i++)
@@ -799,8 +847,8 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
         // Act - Get first page immediately
         var result1 = await store.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
 
-        // Wait for tasks to expire
+        // Advance time past TTL to make tasks expire
-        await Task.Delay(TimeSpan.FromSeconds(1.5), TestContext.Current.CancellationToken);
+        fakeTime.Advance(ttl + TimeSpan.FromMilliseconds(500));
 
         // Get second page after expiration
         var result2 = await store.ListTasksAsync(cursor: result1.NextCursor, cancellationToken: TestContext.Current.CancellationToken);
@@ -1016,16 +1064,26 @@ public class InMemoryMcpTaskStoreTests : LoggedTest
     [Fact]
     public async Task CreateTaskAsync_MaxTasksPerSession_ExcludesExpiredTasks()
     {
-        // Arrange - Short TTL and per-session limit of 1
+        // Arrange - Use FakeTimeProvider for deterministic testing
+        var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
         var shortTtl = TimeSpan.FromMilliseconds(50);
-        using var store = new InMemoryMcpTaskStore(defaultTtl: shortTtl, maxTasksPerSession: 1);
+        using var store = new TestInMemoryMcpTaskStore(
+            defaultTtl: shortTtl,
+            maxTtl: null,
+            pollInterval: null,
+            cleanupInterval: Timeout.InfiniteTimeSpan,
+            pageSize: 100,
+            maxTasks: null,
+            maxTasksPerSession: 1,
+            timeProvider: fakeTime);
+            
         var metadata = new McpTaskMetadata();
 
         // Create first task
         await store.CreateTaskAsync(metadata, new RequestId("req-1"), new JsonRpcRequest { Method = "test" }, "session-1", TestContext.Current.CancellationToken);
 
-        // Wait for it to expire
+        // Advance time past TTL to make the first task expire
-        await Task.Delay(100, TestContext.Current.CancellationToken);
+        fakeTime.Advance(shortTtl + TimeSpan.FromMilliseconds(1));
 
         // Should be able to create another task since the first one expired
         var task2 = await store.CreateTaskAsync(metadata, new RequestId("req-2"), new JsonRpcRequest { Method = "test" }, "session-1", TestContext.Current.CancellationToken);