Add TimeProvider support and test for identical timestamp handling

Co-authored-by: eiriktsarpalis <2813363+eiriktsarpalis@users.noreply.github.com>

Author: Copilot

Directory.Packages.props

@@ -26,6 +26,7 @@
   <!-- Product dependencies .NET Standard -->
   <ItemGroup Condition="'$(TargetFramework)' == 'netstandard2.0'">
     <PackageVersion Include="Microsoft.Bcl.Memory" Version="$(System10Version)" />
+    <PackageVersion Include="Microsoft.Bcl.TimeProvider" Version="$(System10Version)" />
     <PackageVersion Include="System.Diagnostics.DiagnosticSource" Version="$(System10Version)" />
     <PackageVersion Include="System.IO.Pipelines" Version="$(System10Version)" />
     <PackageVersion Include="System.Text.Json" Version="$(System10Version)" />

src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj

@@ -33,6 +33,7 @@
   <ItemGroup Condition="'$(TargetFramework)' == 'netstandard2.0'">
     <Compile Include="..\Common\CancellableStreamReader\**\*.cs" />
     <PackageReference Include="Microsoft.Bcl.Memory" />
+    <PackageReference Include="Microsoft.Bcl.TimeProvider" />
     <PackageReference Include="System.Diagnostics.DiagnosticSource" />
     <PackageReference Include="System.Text.Json" />
     <PackageReference Include="System.Threading.Channels" />

src/ModelContextProtocol.Core/Server/InMemoryMcpTaskStore.cs

@@ -35,6 +35,7 @@ public sealed class InMemoryMcpTaskStore : IMcpTaskStore, IDisposable
     private readonly int _pageSize;
     private readonly int? _maxTasks;
     private readonly int? _maxTasksPerSession;
+    private readonly TimeProvider _timeProvider;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="InMemoryMcpTaskStore"/> class.
@@ -65,14 +66,19 @@ public sealed class InMemoryMcpTaskStore : IMcpTaskStore, IDisposable
     /// Maximum number of tasks allowed per session. Null means unlimited.
     /// When the limit is reached for a session, <see cref="CreateTaskAsync"/> will throw <see cref="InvalidOperationException"/>.
     /// </param>
+    /// <param name="timeProvider">
+    /// Time provider for getting the current time. Defaults to <see cref="TimeProvider.System"/>.
+    /// This parameter is primarily useful for testing scenarios where you need to control time.
+    /// </param>
     public InMemoryMcpTaskStore(
         TimeSpan? defaultTtl = null,
         TimeSpan? maxTtl = null,
         TimeSpan? pollInterval = null,
         TimeSpan? cleanupInterval = null,
         int pageSize = 100,
         int? maxTasks = null,
-        int? maxTasksPerSession = null)
+        int? maxTasksPerSession = null,
+        TimeProvider? timeProvider = null)
     {
         if (defaultTtl.HasValue && maxTtl.HasValue && defaultTtl.Value > maxTtl.Value)
         {
@@ -120,6 +126,7 @@ public InMemoryMcpTaskStore(
         _pageSize = pageSize;
         _maxTasks = maxTasks;
         _maxTasksPerSession = maxTasksPerSession;
+        _timeProvider = timeProvider ?? TimeProvider.System;
 
         cleanupInterval ??= TimeSpan.FromMinutes(1);
         if (cleanupInterval.Value != Timeout.InfiniteTimeSpan)
@@ -155,7 +162,7 @@ public Task<McpTask> CreateTaskAsync(
         }
 
         var taskId = GenerateTaskId();
-        var now = DateTimeOffset.UtcNow;
+        var now = _timeProvider.GetUtcNow();
 
         // Determine TTL: use requested, fall back to default, respect max limit
         var ttl = taskParams.TimeToLive ?? _defaultTtl;
@@ -242,7 +249,7 @@ public Task<McpTask> StoreTaskResultAsync(
             var updatedEntry = new TaskEntry(entry)
             {
                 Status = status,
-                LastUpdatedAt = DateTimeOffset.UtcNow,
+                LastUpdatedAt = _timeProvider.GetUtcNow(),
                 StoredResult = result
             };
 
@@ -303,7 +310,7 @@ public Task<McpTask> UpdateTaskStatusAsync(
             {
                 Status = status,
                 StatusMessage = statusMessage,
-                LastUpdatedAt = DateTimeOffset.UtcNow,
+                LastUpdatedAt = _timeProvider.GetUtcNow(),
             };
 
             if (_tasks.TryUpdate(taskId, updatedEntry, entry))
@@ -398,7 +405,7 @@ public Task<McpTask> CancelTaskAsync(string taskId, string? sessionId = null, Ca
             var updatedEntry = new TaskEntry(entry)
             {
                 Status = McpTaskStatus.Cancelled,
-                LastUpdatedAt = DateTimeOffset.UtcNow,
+                LastUpdatedAt = _timeProvider.GetUtcNow(),
             };
 
             if (_tasks.TryUpdate(taskId, updatedEntry, entry))
@@ -423,15 +430,15 @@ public void Dispose()
     private static bool IsTerminalStatus(McpTaskStatus status) =>
         status is McpTaskStatus.Completed or McpTaskStatus.Failed or McpTaskStatus.Cancelled;
 
-    private static bool IsExpired(TaskEntry entry)
+    private bool IsExpired(TaskEntry entry)
     {
         if (entry.TimeToLive == null)
         {
             return false; // Unlimited lifetime
         }
 
         var expirationTime = entry.CreatedAt + entry.TimeToLive.Value;
-        return DateTimeOffset.UtcNow >= expirationTime;
+        return _timeProvider.GetUtcNow() >= expirationTime;
     }
 
     private void CleanupExpiredTasks(object? state)

tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj

@@ -43,6 +43,7 @@
     <PackageReference Include="Microsoft.Extensions.AI.OpenAI" />
     <PackageReference Include="Microsoft.Extensions.Logging" />
     <PackageReference Include="Microsoft.Extensions.Logging.Console" />
+    <PackageReference Include="Microsoft.Extensions.TimeProvider.Testing" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" />
     <PackageReference Include="Moq" />
     <PackageReference Include="OpenTelemetry" />

tests/ModelContextProtocol.Tests/Server/InMemoryMcpTaskStoreTests.cs

@@ -1,3 +1,4 @@
+using Microsoft.Extensions.Time.Testing;
 using ModelContextProtocol.Protocol;
 using ModelContextProtocol.Server;
 using ModelContextProtocol.Tests.Utils;
@@ -1031,4 +1032,53 @@ public async Task CreateTaskAsync_MaxTasksPerSession_ExcludesExpiredTasks()
         // Assert
         Assert.NotNull(task2);
     }
+
+    [Fact]
+    public async Task ListTasksAsync_KeysetPaginationWorksWithIdenticalTimestamps()
+    {
+        // Arrange - Use a fake time provider to create tasks with identical timestamps
+        var fakeTime = new FakeTimeProvider(DateTimeOffset.UtcNow);
+        using var store = new InMemoryMcpTaskStore(pageSize: 5, timeProvider: fakeTime);
+
+        // Create 10 tasks - all with the EXACT same timestamp
+        var createdTasks = new List<McpTask>();
+        for (int i = 0; i < 10; i++)
+        {
+            var task = await store.CreateTaskAsync(
+                new McpTaskMetadata(),
+                new RequestId($"req-{i}"),
+                new JsonRpcRequest { Method = "test" },
+                null,
+                TestContext.Current.CancellationToken);
+            createdTasks.Add(task);
+        }
+
+        // Verify all tasks have the same CreatedAt timestamp
+        var firstTimestamp = createdTasks[0].CreatedAt;
+        Assert.All(createdTasks, task => Assert.Equal(firstTimestamp, task.CreatedAt));
+
+        // Act - Get first page
+        var result1 = await store.ListTasksAsync(cancellationToken: TestContext.Current.CancellationToken);
+
+        // Assert - First page should have 5 tasks
+        Assert.Equal(5, result1.Tasks.Length);
+        Assert.NotNull(result1.NextCursor);
+
+        // Get second page using cursor
+        var result2 = await store.ListTasksAsync(cursor: result1.NextCursor, cancellationToken: TestContext.Current.CancellationToken);
+
+        // Assert - Second page should have 5 tasks
+        Assert.Equal(5, result2.Tasks.Length);
+        Assert.Null(result2.NextCursor); // No more pages
+
+        // Verify no overlap between pages
+        var page1Ids = result1.Tasks.Select(t => t.TaskId).ToHashSet();
+        var page2Ids = result2.Tasks.Select(t => t.TaskId).ToHashSet();
+        Assert.Empty(page1Ids.Intersect(page2Ids));
+
+        // Verify we got all 10 tasks exactly once
+        var allReturnedIds = page1Ids.Union(page2Ids).ToHashSet();
+        var allCreatedIds = createdTasks.Select(t => t.TaskId).ToHashSet();
+        Assert.Equal(allCreatedIds, allReturnedIds);
+    }
 }