Fix keyset pagination with monotonic UUIDv7-like task IDs (#1215)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: eiriktsarpalis <2813363+eiriktsarpalis@users.noreply.github.com>
Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com>
Co-authored-by: Stephen Toub <stoub@microsoft.com>

Author: 3 people

src/Common/Polyfills/System/IdHelpers.cs

@@ -0,0 +1,40 @@
+using System.Threading;
+
+namespace System;
+
+/// <summary>
+/// Provides helper methods for monotonic ID generation.
+/// </summary>
+internal static class IdHelpers
+{
+    private static long s_counter;
+
+    /// <summary>
+    /// Creates a strictly monotonically increasing identifier string using 64-bit timestamp ticks
+    /// and a 64-bit counter, formatted as a 32-character hexadecimal string (GUID-like).
+    /// </summary>
+    /// <param name="timestamp">The timestamp to embed in the identifier.</param>
+    /// <returns>A new strictly monotonically increasing identifier string.</returns>
+    /// <remarks>
+    /// <para>
+    /// This method creates a 128-bit identifier composed of two 64-bit values:
+    /// - High 64 bits: <see cref="DateTimeOffset.Ticks"/> from the timestamp
+    /// - Low 64 bits: A globally monotonically increasing counter
+    /// </para>
+    /// <para>
+    /// The resulting string is strictly monotonically increasing when compared lexicographically,
+    /// which is required for keyset pagination to work correctly. Unlike <c>Guid.CreateVersion7</c>,
+    /// which uses random bits for intra-millisecond uniqueness, this implementation guarantees
+    /// strict ordering for all identifiers regardless of when they were created.
+    /// </para>
+    /// </remarks>
+    public static string CreateMonotonicId(DateTimeOffset timestamp)
+    {
+        long ticks = timestamp.UtcTicks;
+        long counter = Interlocked.Increment(ref s_counter);
+
+        // Format as 32-character hex string (16 bytes = 128 bits)
+        // High 64 bits: timestamp ticks, Low 64 bits: counter
+        return $"{ticks:x16}{counter:x16}";
+    }
+}

src/ModelContextProtocol.Core/ModelContextProtocol.Core.csproj

@@ -7,7 +7,6 @@
     <PackageId>ModelContextProtocol.Core</PackageId>
     <Description>Core .NET SDK for the Model Context Protocol (MCP)</Description>
     <PackageReadmeFile>README.md</PackageReadmeFile>
-    <AllowUnsafeBlocks>True</AllowUnsafeBlocks>
     <!-- Suppress the experimental tasks warning -->
     <NoWarn>$(NoWarn);MCPEXP001</NoWarn>
   </PropertyGroup>
@@ -19,6 +18,7 @@
   <PropertyGroup Condition="'$(TargetFramework)' == 'netstandard2.0'">
     <!-- CS0436: Allow ObsoleteAttribute to be redefined internally -->
     <NoWarn>$(NoWarn);CS0436</NoWarn>
+    <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
   </PropertyGroup>
 
   <ItemGroup>

src/ModelContextProtocol.Core/Server/InMemoryMcpTaskStore.cs

@@ -3,7 +3,11 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Text.Json;
 
+#if MCP_TEST_TIME_PROVIDER
+namespace ModelContextProtocol.Tests.Internal;
+#else
 namespace ModelContextProtocol;
+#endif
 
 /// <summary>
 /// Provides an in-memory implementation of <see cref="IMcpTaskStore"/> for development and testing.
@@ -35,6 +39,9 @@ public sealed class InMemoryMcpTaskStore : IMcpTaskStore, IDisposable
     private readonly int _pageSize;
     private readonly int? _maxTasks;
     private readonly int? _maxTasksPerSession;
+#if MCP_TEST_TIME_PROVIDER
+    private readonly TimeProvider _timeProvider;
+#endif
 
     /// <summary>
     /// Initializes a new instance of the <see cref="InMemoryMcpTaskStore"/> class.
@@ -120,6 +127,9 @@ public InMemoryMcpTaskStore(
         _pageSize = pageSize;
         _maxTasks = maxTasks;
         _maxTasksPerSession = maxTasksPerSession;
+#if MCP_TEST_TIME_PROVIDER
+        _timeProvider = TimeProvider.System;
+#endif
 
         cleanupInterval ??= TimeSpan.FromMinutes(1);
         if (cleanupInterval.Value != Timeout.InfiniteTimeSpan)
@@ -128,6 +138,26 @@ public InMemoryMcpTaskStore(
         }
     }
 
+#if MCP_TEST_TIME_PROVIDER
+    /// <summary>
+    /// Initializes a new instance of the <see cref="InMemoryMcpTaskStore"/> class with a custom time provider.
+    /// This constructor is only available for testing purposes.
+    /// </summary>
+    internal InMemoryMcpTaskStore(
+        TimeSpan? defaultTtl,
+        TimeSpan? maxTtl,
+        TimeSpan? pollInterval,
+        TimeSpan? cleanupInterval,
+        int pageSize,
+        int? maxTasks,
+        int? maxTasksPerSession,
+        TimeProvider timeProvider)
+        : this(defaultTtl, maxTtl, pollInterval, cleanupInterval, pageSize, maxTasks, maxTasksPerSession)
+    {
+        _timeProvider = timeProvider ?? TimeProvider.System;
+    }
+#endif
+
     /// <inheritdoc/>
     public Task<McpTask> CreateTaskAsync(
         McpTaskMetadata taskParams,
@@ -155,7 +185,7 @@ public Task<McpTask> CreateTaskAsync(
         }
 
         var taskId = GenerateTaskId();
-        var now = DateTimeOffset.UtcNow;
+        var now = GetUtcNow();
 
         // Determine TTL: use requested, fall back to default, respect max limit
         var ttl = taskParams.TimeToLive ?? _defaultTtl;
@@ -242,7 +272,7 @@ public Task<McpTask> StoreTaskResultAsync(
             var updatedEntry = new TaskEntry(entry)
             {
                 Status = status,
-                LastUpdatedAt = DateTimeOffset.UtcNow,
+                LastUpdatedAt = GetUtcNow(),
                 StoredResult = result
             };
 
@@ -303,7 +333,7 @@ public Task<McpTask> UpdateTaskStatusAsync(
             {
                 Status = status,
                 StatusMessage = statusMessage,
-                LastUpdatedAt = DateTimeOffset.UtcNow,
+                LastUpdatedAt = GetUtcNow(),
             };
 
             if (_tasks.TryUpdate(taskId, updatedEntry, entry))
@@ -321,32 +351,22 @@ public Task<ListTasksResult> ListTasksAsync(
         string? sessionId = null,
         CancellationToken cancellationToken = default)
     {
-        // Parse cursor: format is "CreatedAt|TaskId" for keyset pagination
-        (DateTimeOffset, string)? parsedCursor = null;
-        if (cursor != null)
-        {
-            var parts = cursor.Split('|');
-            if (parts.Length == 2 &&
-                DateTimeOffset.TryParse(parts[0], out var parsedDate))
-            {
-                parsedCursor = (parsedDate, parts[1]);
-            }
-        }
-
         // Stream enumeration - filter by session, exclude expired, apply keyset pagination
         var query = _tasks.Values
             .Where(e => sessionId == null || e.SessionId == sessionId)
             .Where(e => !IsExpired(e));
 
-        // Apply keyset filter if cursor provided: (CreatedAt, TaskId) > cursor
-        if (parsedCursor is { } parsedCursorValue)
+        // Apply keyset filter if cursor provided: TaskId > cursor
+        // UUID v7 task IDs are monotonically increasing and inherently time-ordered
+        if (cursor != null)
         {
-            query = query.Where(e => (e.CreatedAt, e.TaskId).CompareTo(parsedCursorValue) > 0);
+            query = query.Where(e => string.CompareOrdinal(e.TaskId, cursor) > 0);
         }
 
-        // Order by (CreatedAt, TaskId) for stable, deterministic pagination
+        // Order by TaskId for stable, deterministic pagination
+        // UUID v7 task IDs sort chronologically due to embedded timestamp
         var page = query
-            .OrderBy(e => (e.CreatedAt, e.TaskId))
+            .OrderBy(e => e.TaskId, StringComparer.Ordinal)
             .Take(_pageSize + 1) // Take one extra to check if there's a next page
             .Select(e => e.ToMcpTask())
             .ToList();
@@ -356,7 +376,7 @@ public Task<ListTasksResult> ListTasksAsync(
         if (page.Count > _pageSize)
         {
             var lastItemInPage = page[_pageSize - 1]; // Last item we'll actually return
-            nextCursor = $"{lastItemInPage.CreatedAt:O}|{lastItemInPage.TaskId}";
+            nextCursor = lastItemInPage.TaskId;
             page.RemoveAt(_pageSize); // Remove the extra item
         }
         else
@@ -397,7 +417,7 @@ public Task<McpTask> CancelTaskAsync(string taskId, string? sessionId = null, Ca
             var updatedEntry = new TaskEntry(entry)
             {
                 Status = McpTaskStatus.Cancelled,
-                LastUpdatedAt = DateTimeOffset.UtcNow,
+                LastUpdatedAt = GetUtcNow(),
             };
 
             if (_tasks.TryUpdate(taskId, updatedEntry, entry))
@@ -417,20 +437,31 @@ public void Dispose()
         _cleanupTimer?.Dispose();
     }
 
-    private static string GenerateTaskId() => Guid.NewGuid().ToString("N");
+    private string GenerateTaskId() =>
+        IdHelpers.CreateMonotonicId(GetUtcNow());
 
     private static bool IsTerminalStatus(McpTaskStatus status) =>
         status is McpTaskStatus.Completed or McpTaskStatus.Failed or McpTaskStatus.Cancelled;
 
+#if MCP_TEST_TIME_PROVIDER
+    private DateTimeOffset GetUtcNow() => _timeProvider.GetUtcNow();
+#else
+    private static DateTimeOffset GetUtcNow() => DateTimeOffset.UtcNow;
+#endif
+
+#if MCP_TEST_TIME_PROVIDER
+    private bool IsExpired(TaskEntry entry)
+#else
     private static bool IsExpired(TaskEntry entry)
+#endif
     {
         if (entry.TimeToLive == null)
         {
             return false; // Unlimited lifetime
         }
 
         var expirationTime = entry.CreatedAt + entry.TimeToLive.Value;
-        return DateTimeOffset.UtcNow >= expirationTime;
+        return GetUtcNow() >= expirationTime;
     }
 
     private void CleanupExpiredTasks(object? state)

tests/ModelContextProtocol.Tests/ModelContextProtocol.Tests.csproj

@@ -5,12 +5,15 @@
     <TargetFrameworks>net10.0;net9.0;net8.0;net472</TargetFrameworks>
     <ImplicitUsings>enable</ImplicitUsings>
     <Nullable>enable</Nullable>
+    <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
 
     <IsPackable>false</IsPackable>
     <IsTestProject>true</IsTestProject>
     <RootNamespace>ModelContextProtocol.Tests</RootNamespace>
     <!-- https://github.com/dotnet/sdk/issues/51060 -->
     <NoWarn>$(NoWarn);NU1903;NU1902</NoWarn>
+    <!-- Define for conditional TimeProvider support in InMemoryMcpTaskStore -->
+    <DefineConstants>$(DefineConstants);MCP_TEST_TIME_PROVIDER</DefineConstants>
   </PropertyGroup>
 
   <PropertyGroup Condition="'$(TargetFramework)' == 'net9.0'">
@@ -27,13 +30,21 @@
 
   <ItemGroup>
     <Compile Include="..\Common\**\*.cs" />
+    <!-- Link InMemoryMcpTaskStore.cs for testing with TimeProvider support -->
+    <Compile Include="..\..\src\ModelContextProtocol.Core\Server\InMemoryMcpTaskStore.cs" Link="Server\InMemoryMcpTaskStore.cs" />
+    <!-- Include dependencies for the linked InMemoryMcpTaskStore.cs -->
+    <Compile Include="..\..\src\Common\Experimentals.cs" Link="Experimentals.cs" />
   </ItemGroup>
 
   <ItemGroup Condition="'$(TargetFramework)' == 'net472'">
     <Compile Include="..\..\src\Common\Throw.cs" Link="Throw.cs" />
     <Compile Include="..\..\src\Common\Polyfills\**\*.cs" />
   </ItemGroup>
 
+  <ItemGroup Condition="'$(TargetFramework)' != 'net472'">
+    <Compile Include="..\..\src\Common\Polyfills\System\IdHelpers.cs" Link="Polyfills\IdHelpers.cs" />
+  </ItemGroup>
+
   <ItemGroup>
     <PackageReference Include="coverlet.collector">
       <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
@@ -43,6 +54,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" />