diff --git a/src/ModelContextProtocol.Core/Server/InMemoryMcpTaskStore.cs b/src/ModelContextProtocol.Core/Server/InMemoryMcpTaskStore.cs index 27156e98d..b2f9b050d 100644 --- a/src/ModelContextProtocol.Core/Server/InMemoryMcpTaskStore.cs +++ b/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 InMemoryMcpTaskStore( 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 } } diff --git a/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs b/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs index 5b9db455a..56e9aaca9 100644 --- a/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs +++ b/tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs @@ -235,14 +235,24 @@ public async Task UpdateTaskStatusAsync_UpdatesStatus() [Fact] public async Task UpdateTaskStatusAsync_UpdatesLastUpdatedAt() { - // Arrange - using var store = new InMemoryMcpTaskStore(); + // Arrange - Use FakeTimeProvider for deterministic testing + 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 - await Task.Delay(10, TestContext.Current.CancellationToken); + // Advance time to ensure timestamp changes + fakeTime.Advance(TimeSpan.FromMilliseconds(10)); // Act await store.UpdateTaskStatusAsync(task.TaskId, McpTaskStatus.Working, null, null, TestContext.Current.CancellationToken); @@ -460,16 +470,28 @@ await Assert.ThrowsAsync( [Fact] public async Task Dispose_StopsCleanupTimer() { - // Arrange - var store = new InMemoryMcpTaskStore(cleanupInterval: TimeSpan.FromMilliseconds(100)); - var metadata = new McpTaskMetadata { TimeToLive = TimeSpan.FromMilliseconds(100) }; // Very short TTL + // Arrange - Use FakeTimeProvider for deterministic testing + var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow); + 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 - await Task.Delay(300, TestContext.Current.CancellationToken); + // Advance time - timer should not fire after dispose + 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 async Task Dispose_StopsCleanupTimer() [Fact] public async Task CleanupExpiredTasks_RemovesExpiredTasks() { - // Arrange - using var store = new InMemoryMcpTaskStore(cleanupInterval: TimeSpan.FromMilliseconds(50)); - var metadata = new McpTaskMetadata { TimeToLive = TimeSpan.FromMilliseconds(100) }; // 100ms TTL + // Arrange - Use FakeTimeProvider for deterministic testing + var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow); + 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) - await Task.Delay(250, TestContext.Current.CancellationToken); + // Advance time past the TTL to make task expired + 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 async Task ListTasksAsync_NoDuplicatesWithIdenticalTimestamps() [Fact] public async Task ListTasksAsync_ConsistentWithExpiredTasksRemovedBetweenPages() { - // Arrange - Use TTL of 1 second - using var store = new InMemoryMcpTaskStore(defaultTtl: TimeSpan.FromSeconds(1), pageSize: 5, cleanupInterval: Timeout.InfiniteTimeSpan); + // Arrange - Use FakeTimeProvider for deterministic testing + 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 async Task ListTasksAsync_ConsistentWithExpiredTasksRemovedBetweenPages() // Act - Get first page immediately var result1 = await store.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken); - // Wait for tasks to expire - await Task.Delay(TimeSpan.FromSeconds(1.5), TestContext.Current.CancellationToken); + // Advance time past TTL to make tasks expire + 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 @@ await Assert.ThrowsAsync(() => [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 - await Task.Delay(100, TestContext.Current.CancellationToken); + // Advance time past TTL to make the first task expire + 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);