Fix pagination handling in McpServer (#177)

- We were adding tools/prompts from the collections on every request. If multiple requests came in with different cursors, we'd re-add the same tools each time.
- We were defeating the purpose of pagination by doing all of the aggregation in the server. If a custom handler returns a paginated result, we should instead propagate that back to the client, who can choose to get more results when needed.

src/ModelContextProtocol/Server/AIFunctionMcpServerTool.cs

@@ -1,7 +1,6 @@
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.DependencyInjection;
 using ModelContextProtocol.Protocol.Types;
-using ModelContextProtocol.Shared;
 using ModelContextProtocol.Utils;
 using ModelContextProtocol.Utils.Json;
 using System.Diagnostics.CodeAnalysis;

src/ModelContextProtocol/Server/McpServer.cs

@@ -1,5 +1,4 @@
 using Microsoft.Extensions.Logging;
-using ModelContextProtocol.Logging;
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
@@ -214,41 +213,26 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
             throw new McpServerException("ListPrompts and GetPrompt handlers should be specified together.");
         }
 
-        // Handle tools provided via DI.
+        // Handle prompts provided via DI.
         if (prompts is { IsEmpty: false })
         {
+            // Synthesize the handlers, making sure a PromptsCapability is specified.
             var originalListPromptsHandler = listPromptsHandler;
-            var originalGetPromptHandler = getPromptHandler;
-
-            // Synthesize the handlers, making sure a ToolsCapability is specified.
             listPromptsHandler = async (request, cancellationToken) =>
             {
-                ListPromptsResult result = new();
+                ListPromptsResult result = originalListPromptsHandler is not null ?
-                foreach (McpServerPrompt prompt in prompts)
+                    await originalListPromptsHandler(request, cancellationToken).ConfigureAwait(false) :
-                {
+                    new();
-                    result.Prompts.Add(prompt.ProtocolPrompt);
-                }
 
-                if (originalListPromptsHandler is not null)
+                if (request.Params?.Cursor is null)
                 {
-                    string? nextCursor = null;
+                    result.Prompts.AddRange(prompts.Select(t => t.ProtocolPrompt));
-                    do
-                    {
-                        ListPromptsResult extraResults = await originalListPromptsHandler(request, cancellationToken).ConfigureAwait(false);
-                        result.Prompts.AddRange(extraResults.Prompts);
-
-                        nextCursor = extraResults.NextCursor;
-                        if (nextCursor is not null)
-                        {
-                            request = request with { Params = new() { Cursor = nextCursor } };
-                        }
-                    }
-                    while (nextCursor is not null);
                 }
 
                 return result;
             };
 
+            var originalGetPromptHandler = getPromptHandler;
             getPromptHandler = (request, cancellationToken) =>
             {
                 if (request.Params is null ||
@@ -316,38 +300,23 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
         // Handle tools provided via DI.
         if (tools is { IsEmpty: false })
         {
-            var originalListToolsHandler = listToolsHandler;
-            var originalCallToolHandler = callToolHandler;
-
             // Synthesize the handlers, making sure a ToolsCapability is specified.
+            var originalListToolsHandler = listToolsHandler;
             listToolsHandler = async (request, cancellationToken) =>
             {
-                ListToolsResult result = new();
+                ListToolsResult result = originalListToolsHandler is not null ?
-                foreach (McpServerTool tool in tools)
+                    await originalListToolsHandler(request, cancellationToken).ConfigureAwait(false) :
-                {
+                    new();
-                    result.Tools.Add(tool.ProtocolTool);
-                }
 
-                if (originalListToolsHandler is not null)
+                if (request.Params?.Cursor is null)
                 {
-                    string? nextCursor = null;
+                    result.Tools.AddRange(tools.Select(t => t.ProtocolTool));
-                    do
-                    {
-                        ListToolsResult extraResults = await originalListToolsHandler(request, cancellationToken).ConfigureAwait(false);
-                        result.Tools.AddRange(extraResults.Tools);
-
-                        nextCursor = extraResults.NextCursor;
-                        if (nextCursor is not null)
-                        {
-                            request = request with { Params = new() { Cursor = nextCursor } };
-                        }
-                    }
-                    while (nextCursor is not null);
                 }
 
                 return result;
             };
 
+            var originalCallToolHandler = callToolHandler;
             callToolHandler = (request, cancellationToken) =>
             {
                 if (request.Params is null ||

src/ModelContextProtocol/TokenProgress.cs

@@ -1,6 +1,5 @@
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Server;
-using ModelContextProtocol.Shared;
 
 namespace ModelContextProtocol;
 

tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs

@@ -2,7 +2,6 @@ using Microsoft.Extensions.DependencyInjection;
 using ModelContextProtocol.Client;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Server;
-using ModelContextProtocol.Tests.Transport;
 using ModelContextProtocol.Tests.Utils;
 using System.IO.Pipelines;
 

tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsPromptsTests.cs

@@ -6,12 +6,13 @@ using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
-using ModelContextProtocol.Tests.Transport;
 using ModelContextProtocol.Tests.Utils;
 using System.ComponentModel;
 using System.IO.Pipelines;
 using System.Threading.Channels;
 
+#pragma warning disable CS1998 // Async method lacks 'await' operators and will run synchronously
+
 namespace ModelContextProtocol.Tests.Configuration;
 
 public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposable
@@ -28,7 +29,70 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
     {
         ServiceCollection sc = new();
         sc.AddSingleton(LoggerFactory);
-        _builder = sc.AddMcpServer().WithStdioServerTransport().WithPrompts<SimplePrompts>();
+        _builder = sc
+            .AddMcpServer()
+            .WithStdioServerTransport()
+            .WithListPromptsHandler(async (request, cancellationToken) =>
+            {
+                var cursor = request.Params?.Cursor;
+                switch (cursor)
+                {
+                    case null:
+                        return new()
+                        {
+                            NextCursor = "abc",
+                            Prompts = [new()
+                            {
+                                Name = "FirstCustomPrompt",
+                                Description = "First prompt returned by custom handler",
+                            }],
+                        };
+
+                    case "abc":
+                        return new()
+                        {
+                            NextCursor = "def",
+                            Prompts = [new()
+                            {
+                                Name = "SecondCustomPrompt",
+                                Description = "Second prompt returned by custom handler",
+                            }],
+                        };
+
+                    case "def":
+                        return new()
+                        {
+                            NextCursor = null,
+                            Prompts = [new()
+                            {
+                                Name = "FinalCustomPrompt",
+                                Description = "Final prompt returned by custom handler",
+                            }],
+                        };
+
+                    default:
+                        throw new Exception("Unexpected cursor");
+                }
+            })
+            .WithGetPromptHandler(async (request, cancellationToken) =>
+            {
+                switch (request.Params?.Name)
+                {
+                    case "FirstCustomPrompt":
+                    case "SecondCustomPrompt":
+                    case "FinalCustomPrompt":
+                        return new GetPromptResult()
+                        {
+                            Messages = [new() { Role = Role.User, Content = new() { Text = $"hello from {request.Params.Name}", Type = "text" } }],
+                        };
+
+                    default:
+                        throw new Exception($"Unknown prompt '{request.Params?.Name}'");
+                }
+            })
+            .WithPrompts<SimplePrompts>();
+
+
         // Call WithStdioServerTransport to get the IMcpServer registration, then overwrite default transport with a pipe transport.
         sc.AddSingleton<ITransport>(new StreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream(), loggerFactory: LoggerFactory));
         sc.AddSingleton(new ObjectWithId());
@@ -85,7 +149,7 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
         IMcpClient client = await CreateMcpClientForServer();
 
         var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(3, prompts.Count);
+        Assert.Equal(6, prompts.Count);
 
         var prompt = prompts.First(t => t.Name == nameof(SimplePrompts.ReturnsChatMessages));
         Assert.Equal("Returns chat messages", prompt.Description);
@@ -98,6 +162,14 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
         Assert.Equal(2, chatMessages.Count);
         Assert.Equal("The prompt is: hello", chatMessages[0].Text);
         Assert.Equal("Summarize.", chatMessages[1].Text);
+
+        prompt = prompts.First(t => t.Name == "SecondCustomPrompt");
+        Assert.Equal("Second prompt returned by custom handler", prompt.Description);
+        result = await prompt.GetAsync(cancellationToken: TestContext.Current.CancellationToken);
+        chatMessages = result.ToChatMessages();
+        Assert.NotNull(chatMessages);
+        Assert.Single(chatMessages);
+        Assert.Equal("hello from SecondCustomPrompt", chatMessages[0].Text);
     }
 
     [Fact]
@@ -106,7 +178,7 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
         IMcpClient client = await CreateMcpClientForServer();
 
         var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(3, prompts.Count);
+        Assert.Equal(6, prompts.Count);
 
         Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
         client.AddNotificationHandler("notifications/prompts/list_changed", notification =>
@@ -127,7 +199,7 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
         await notificationRead;
 
         prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(4, prompts.Count);
+        Assert.Equal(7, prompts.Count);
         Assert.Contains(prompts, t => t.Name == "NewPrompt");
 
         notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
@@ -136,7 +208,7 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
         await notificationRead;
 
         prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(3, prompts.Count);
+        Assert.Equal(6, prompts.Count);
         Assert.DoesNotContain(prompts, t => t.Name == "NewPrompt");
     }
 

tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs

@@ -7,7 +7,6 @@ using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;
 using ModelContextProtocol.Server;
-using ModelContextProtocol.Tests.Transport;
 using ModelContextProtocol.Tests.Utils;
 using System.Collections.Concurrent;
 using System.ComponentModel;
@@ -16,6 +15,8 @@ using System.Text.Json;
 using System.Text.RegularExpressions;
 using System.Threading.Channels;
 
+#pragma warning disable CS1998 // Async method lacks 'await' operators and will run synchronously
+
 namespace ModelContextProtocol.Tests.Configuration;
 
 public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
@@ -32,7 +33,90 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
     {
         ServiceCollection sc = new();
         sc.AddSingleton(LoggerFactory);
-        _builder = sc.AddMcpServer().WithStdioServerTransport().WithTools<EchoTool>();
+        _builder = sc
+            .AddMcpServer()
+            .WithStdioServerTransport()
+            .WithListToolsHandler(async (request, cancellationToken) =>
+            {
+                var cursor = request.Params?.Cursor;
+                switch (cursor)
+                {
+                    case null:
+                        return new()
+                        {
+                            NextCursor = "abc",
+                            Tools = [new()
+                            {
+                                Name = "FirstCustomTool",
+                                Description = "First tool returned by custom handler",
+                                InputSchema = JsonSerializer.Deserialize<JsonElement>("""
+                                    {
+                                      "type": "object",
+                                      "properties": {},
+                                      "required": []
+                                    }
+                                    """),
+                            }],
+                        };
+
+                    case "abc":
+                        return new()
+                        {
+                            NextCursor = "def",
+                            Tools = [new()
+                            {
+                                Name = "SecondCustomTool",
+                                Description = "Second tool returned by custom handler",
+                                InputSchema = JsonSerializer.Deserialize<JsonElement>("""
+                                    {
+                                      "type": "object",
+                                      "properties": {},
+                                      "required": []
+                                    }
+                                    """),
+                            }],
+                        };
+
+                    case "def":
+                        return new()
+                        {
+                            NextCursor = null,
+                            Tools = [new()
+                            {
+                                Name = "FinalCustomTool",
+                                Description = "Third tool returned by custom handler",
+                                InputSchema = JsonSerializer.Deserialize<JsonElement>("""
+                                    {
+                                      "type": "object",
+                                      "properties": {},
+                                      "required": []
+                                    }
+                                    """),
+                            }],
+                        };
+
+                    default:
+                        throw new Exception("Unexpected cursor");
+                }
+            })
+            .WithCallToolHandler(async (request, cancellationToken) =>
+            {
+                switch (request.Params?.Name)
+                {
+                    case "FirstCustomTool":
+                    case "SecondCustomTool":
+                    case "FinalCustomTool":
+                        return new CallToolResponse()
+                        {
+                            Content = [new Content() { Text = $"{request.Params.Name}Result", Type = "text" }],
+                        };
+
+                    default:
+                        throw new Exception($"Unknown tool '{request.Params?.Name}'");
+                }
+            })
+            .WithTools<EchoTool>();
+
         // Call WithStdioServerTransport to get the IMcpServer registration, then overwrite default transport with a pipe transport.
         sc.AddSingleton<ITransport>(new StreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream(), loggerFactory: LoggerFactory));
         sc.AddSingleton(new ObjectWithId());
@@ -89,7 +173,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
         IMcpClient client = await CreateMcpClientForServer();
 
         var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(13, tools.Count);
+        Assert.Equal(16, tools.Count);
 
         McpClientTool echoTool = tools.First(t => t.Name == "Echo");
         Assert.Equal("Echo", echoTool.Name);
@@ -134,7 +218,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
                 cancellationToken: TestContext.Current.CancellationToken))
             {
                 var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-                Assert.Equal(13, tools.Count);
+                Assert.Equal(16, tools.Count);
 
                 McpClientTool echoTool = tools.First(t => t.Name == "Echo");
                 Assert.Equal("Echo", echoTool.Name);
@@ -161,7 +245,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
         IMcpClient client = await CreateMcpClientForServer();
 
         var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(13, tools.Count);
+        Assert.Equal(16, tools.Count);
 
         Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
         client.AddNotificationHandler(NotificationMethods.ToolListChangedNotification, notification =>
@@ -182,7 +266,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
         await notificationRead;
 
         tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(14, tools.Count);
+        Assert.Equal(17, tools.Count);
         Assert.Contains(tools, t => t.Name == "NewTool");
 
         notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
@@ -191,7 +275,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
         await notificationRead;
 
         tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
-        Assert.Equal(13, tools.Count);
+        Assert.Equal(16, tools.Count);
         Assert.DoesNotContain(tools, t => t.Name == "NewTool");
     }
 
@@ -225,9 +309,16 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
 
         Assert.NotNull(result.Content);
         Assert.NotEmpty(result.Content);
-
         Assert.Equal("hello Peter", result.Content[0].Text);
         Assert.Equal("hello2 Peter", result.Content[1].Text);
+
+        result = await client.CallToolAsync(
+            "SecondCustomTool",
+            cancellationToken: TestContext.Current.CancellationToken);
+        Assert.NotNull(result);
+        Assert.NotNull(result.Content);
+        Assert.NotEmpty(result.Content);
+        Assert.Equal("SecondCustomToolResult", result.Content[0].Text);
     }
 
     [Fact]

tests/ModelContextProtocol.Tests/Server/McpServerTests.cs

@@ -1,6 +1,5 @@
 using Microsoft.Extensions.AI;
 using Microsoft.Extensions.DependencyInjection;
-using ModelContextProtocol.Client;
 using ModelContextProtocol.Protocol.Messages;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Protocol.Types;

tests/ModelContextProtocol.Tests/SseServerIntegrationTestFixture.cs

@@ -1,5 +1,4 @@
-using Microsoft.Extensions.Logging;
+using ModelContextProtocol.Client;
-using ModelContextProtocol.Client;
 using ModelContextProtocol.Protocol.Transport;
 using ModelContextProtocol.Test.Utils;
 using ModelContextProtocol.Tests.Utils;