modelcontextprotocol/csharp-sdk
Watch
1
0
Fork
You've already forked csharp-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Fix pagination handling in McpServer (#177)

Browse source 
- 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.
This commit is contained in:
Stephen Toub 2025-04-01 09:54:19 -04:00 • committed by GitHub
commit b12d728700
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 193 additions and 66 deletions
Download patch file Download diff file Expand all files Collapse all files

1
src/ModelContextProtocol/Server/AIFunctionMcpServerTool.cs
View file

@ -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;

61
src/ModelContextProtocol/Server/McpServer.cs
View file

@ -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();
foreach (McpServerPrompt prompt in prompts)
{
result.Prompts.Add(prompt.ProtocolPrompt);
}
ListPromptsResult result = originalListPromptsHandler is not null ?
await originalListPromptsHandler(request, cancellationToken).ConfigureAwait(false) :
new();
if (originalListPromptsHandler is not null)
if (request.Params?.Cursor is null)
{
string? nextCursor = null;
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);
result.Prompts.AddRange(prompts.Select(t => t.ProtocolPrompt));
}
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();
foreach (McpServerTool tool in tools)
{
result.Tools.Add(tool.ProtocolTool);
}
ListToolsResult result = originalListToolsHandler is not null ?
await originalListToolsHandler(request, cancellationToken).ConfigureAwait(false) :
new();
if (originalListToolsHandler is not null)
if (request.Params?.Cursor is null)
{
string? nextCursor = null;
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);
result.Tools.AddRange(tools.Select(t => t.ProtocolTool));
}
return result;
};
var originalCallToolHandler = callToolHandler;
callToolHandler = (request, cancellationToken) =>
{
if (request.Params is null ||

1
src/ModelContextProtocol/TokenProgress.cs
View file

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

1
tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs
View file

@ -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;

84
tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsPromptsTests.cs
View file

@ -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");
}

107
tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs
View file

@ -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]

1
tests/ModelContextProtocol.Tests/Server/McpServerTests.cs
View file

@ -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;

3
tests/ModelContextProtocol.Tests/SseServerIntegrationTestFixture.cs
View file

@ -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;