Skip to content

Use TimeProvider in InMemoryMcpTaskStore for deterministic testing #1303

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Copilot wants to merge 2 commits into
base: main
Choose a base branch
from
+88 −22
Draft

Use TimeProvider in InMemoryMcpTaskStore for deterministic testing #1303

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
54c2b93
Initial plan
Copilot Feb 18, 2026
9dcd132
Update InMemoryMcpTaskStore and tests to use TimeProvider
Copilot Feb 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/ModelContextProtocol.Core/Server/InMemoryMcpTaskStore.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
}
}

Expand Down
102 changes: 80 additions & 22 deletions tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -460,16 +470,28 @@ await Assert.ThrowsAsync<InvalidOperationException>(
[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(cleanupInterval * 3);

// Assert - Store should still be accessible after dispose (no exceptions)
// The cleanup timer should have stopped
Expand All @@ -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);
Expand Down Expand Up @@ -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++)
Expand All @@ -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);
Expand Down Expand Up @@ -1016,16 +1064,26 @@ await Assert.ThrowsAsync<InvalidOperationException>(() =>
[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);
Expand Down