Fix and enhance cancellation operations across MCP Sessions. (#179)
* Propagate CancellationToken request cancellation to remote endpoint * Refactor existing tests to have a shared base class * Add cancellation tests --------- Co-authored-by: Stephen Toub <stoub@microsoft.com>
This commit is contained in:
Tyler Kendrick
GitHub
parent
4dd2f42b82
commit
03e3094dcc
2 changed files with 334 additions and 336 deletions
|
|
@ -3,133 +3,92 @@ using Microsoft.Extensions.DependencyInjection;
|
|||
using Microsoft.Extensions.Options;
|
||||
using ModelContextProtocol.Client;
|
||||
using ModelContextProtocol.Protocol.Messages;
|
||||
using ModelContextProtocol.Protocol.Transport;
|
||||
using ModelContextProtocol.Protocol.Types;
|
||||
using ModelContextProtocol.Server;
|
||||
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
|
||||
public class McpServerBuilderExtensionsPromptsTests : ClientServerTestBase
|
||||
{
|
||||
private readonly Pipe _clientToServerPipe = new();
|
||||
private readonly Pipe _serverToClientPipe = new();
|
||||
private readonly ServiceProvider _serviceProvider;
|
||||
private readonly IMcpServerBuilder _builder;
|
||||
private readonly CancellationTokenSource _cts;
|
||||
private readonly Task _serverTask;
|
||||
|
||||
public McpServerBuilderExtensionsPromptsTests(ITestOutputHelper testOutputHelper)
|
||||
: base(testOutputHelper)
|
||||
{
|
||||
ServiceCollection sc = new();
|
||||
sc.AddSingleton(LoggerFactory);
|
||||
_builder = sc
|
||||
.AddMcpServer()
|
||||
.WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream())
|
||||
.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>();
|
||||
|
||||
|
||||
sc.AddSingleton(new ObjectWithId());
|
||||
_serviceProvider = sc.BuildServiceProvider();
|
||||
|
||||
var server = _serviceProvider.GetRequiredService<IMcpServer>();
|
||||
_cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
|
||||
_serverTask = server.RunAsync(cancellationToken: _cts.Token);
|
||||
}
|
||||
|
||||
public async ValueTask DisposeAsync()
|
||||
protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
|
||||
{
|
||||
await _cts.CancelAsync();
|
||||
mcpServerBuilder
|
||||
.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",
|
||||
}],
|
||||
};
|
||||
|
||||
_clientToServerPipe.Writer.Complete();
|
||||
_serverToClientPipe.Writer.Complete();
|
||||
case "abc":
|
||||
return new()
|
||||
{
|
||||
NextCursor = "def",
|
||||
Prompts = [new()
|
||||
{
|
||||
Name = "SecondCustomPrompt",
|
||||
Description = "Second prompt returned by custom handler",
|
||||
}],
|
||||
};
|
||||
|
||||
await _serverTask;
|
||||
case "def":
|
||||
return new()
|
||||
{
|
||||
NextCursor = null,
|
||||
Prompts = [new()
|
||||
{
|
||||
Name = "FinalCustomPrompt",
|
||||
Description = "Final prompt returned by custom handler",
|
||||
}],
|
||||
};
|
||||
|
||||
await _serviceProvider.DisposeAsync();
|
||||
_cts.Dispose();
|
||||
Dispose();
|
||||
}
|
||||
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" } }],
|
||||
};
|
||||
|
||||
private async Task<IMcpClient> CreateMcpClientForServer(McpClientOptions? options = null)
|
||||
{
|
||||
return await McpClientFactory.CreateAsync(
|
||||
new StreamClientTransport(
|
||||
serverInput: _clientToServerPipe.Writer.AsStream(),
|
||||
serverOutput: _serverToClientPipe.Reader.AsStream(),
|
||||
LoggerFactory),
|
||||
loggerFactory: LoggerFactory,
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
default:
|
||||
throw new Exception($"Unknown prompt '{request.Params?.Name}'");
|
||||
}
|
||||
})
|
||||
.WithPrompts<SimplePrompts>();
|
||||
|
||||
services.AddSingleton(new ObjectWithId());
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public void Adds_Prompts_To_Server()
|
||||
{
|
||||
var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var serverOptions = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var prompts = serverOptions?.Capabilities?.Prompts?.PromptCollection;
|
||||
Assert.NotNull(prompts);
|
||||
Assert.NotEmpty(prompts);
|
||||
|
|
@ -176,7 +135,7 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
|
|||
var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
|
||||
Assert.False(notificationRead.IsCompleted);
|
||||
|
||||
var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var serverOptions = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var serverPrompts = serverOptions.Capabilities?.Prompts?.PromptCollection;
|
||||
Assert.NotNull(serverPrompts);
|
||||
|
||||
|
|
@ -242,7 +201,9 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
|
|||
[Fact]
|
||||
public void WithPrompts_InvalidArgs_Throws()
|
||||
{
|
||||
Assert.Throws<ArgumentNullException>("promptTypes", () => _builder.WithPrompts((IEnumerable<Type>)null!));
|
||||
IMcpServerBuilder builder = new ServiceCollection().AddMcpServer();
|
||||
|
||||
Assert.Throws<ArgumentNullException>("promptTypes", () => builder.WithPrompts((IEnumerable<Type>)null!));
|
||||
|
||||
IMcpServerBuilder nullBuilder = null!;
|
||||
Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithPrompts<object>());
|
||||
|
|
@ -253,9 +214,11 @@ public class McpServerBuilderExtensionsPromptsTests : LoggedTest, IAsyncDisposab
|
|||
[Fact]
|
||||
public void Empty_Enumerables_Is_Allowed()
|
||||
{
|
||||
_builder.WithPrompts(promptTypes: []); // no exception
|
||||
_builder.WithPrompts<object>(); // no exception even though no prompts exposed
|
||||
_builder.WithPromptsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no prompts exposed
|
||||
IMcpServerBuilder builder = new ServiceCollection().AddMcpServer();
|
||||
|
||||
builder.WithPrompts(promptTypes: []); // no exception
|
||||
builder.WithPrompts<object>(); // no exception even though no prompts exposed
|
||||
builder.WithPromptsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no prompts exposed
|
||||
}
|
||||
|
||||
[Fact]
|
||||
|
|
|
|||
|
|
@ -7,7 +7,6 @@ using ModelContextProtocol.Protocol.Messages;
|
|||
using ModelContextProtocol.Protocol.Transport;
|
||||
using ModelContextProtocol.Protocol.Types;
|
||||
using ModelContextProtocol.Server;
|
||||
using ModelContextProtocol.Tests.Utils;
|
||||
using System.Collections.Concurrent;
|
||||
using System.ComponentModel;
|
||||
using System.IO.Pipelines;
|
||||
|
|
@ -19,23 +18,16 @@ using System.Threading.Channels;
|
|||
|
||||
namespace ModelContextProtocol.Tests.Configuration;
|
||||
|
||||
public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
||||
public class McpServerBuilderExtensionsToolsTests : ClientServerTestBase
|
||||
{
|
||||
private readonly Pipe _clientToServerPipe = new();
|
||||
private readonly Pipe _serverToClientPipe = new();
|
||||
private readonly ServiceProvider _serviceProvider;
|
||||
private readonly IMcpServerBuilder _builder;
|
||||
private readonly CancellationTokenSource _cts;
|
||||
private readonly Task _serverTask;
|
||||
|
||||
public McpServerBuilderExtensionsToolsTests(ITestOutputHelper testOutputHelper)
|
||||
: base(testOutputHelper)
|
||||
{
|
||||
ServiceCollection sc = new();
|
||||
sc.AddSingleton(LoggerFactory);
|
||||
_builder = sc
|
||||
.AddMcpServer()
|
||||
.WithStreamServerTransport(_clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream())
|
||||
}
|
||||
|
||||
protected override void ConfigureServices(ServiceCollection services, IMcpServerBuilder mcpServerBuilder)
|
||||
{
|
||||
mcpServerBuilder
|
||||
.WithListToolsHandler(async (request, cancellationToken) =>
|
||||
{
|
||||
var cursor = request.Params?.Cursor;
|
||||
|
|
@ -46,17 +38,17 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
|||
{
|
||||
NextCursor = "abc",
|
||||
Tools = [new()
|
||||
{
|
||||
Name = "FirstCustomTool",
|
||||
Description = "First tool returned by custom handler",
|
||||
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
"""),
|
||||
}],
|
||||
{
|
||||
Name = "FirstCustomTool",
|
||||
Description = "First tool returned by custom handler",
|
||||
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
"""),
|
||||
}],
|
||||
};
|
||||
|
||||
case "abc":
|
||||
|
|
@ -64,17 +56,17 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
|||
{
|
||||
NextCursor = "def",
|
||||
Tools = [new()
|
||||
{
|
||||
Name = "SecondCustomTool",
|
||||
Description = "Second tool returned by custom handler",
|
||||
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
"""),
|
||||
}],
|
||||
{
|
||||
Name = "SecondCustomTool",
|
||||
Description = "Second tool returned by custom handler",
|
||||
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
"""),
|
||||
}],
|
||||
};
|
||||
|
||||
case "def":
|
||||
|
|
@ -82,17 +74,17 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
|||
{
|
||||
NextCursor = null,
|
||||
Tools = [new()
|
||||
{
|
||||
Name = "FinalCustomTool",
|
||||
Description = "Third tool returned by custom handler",
|
||||
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
"""),
|
||||
}],
|
||||
{
|
||||
Name = "FinalCustomTool",
|
||||
Description = "Third tool returned by custom handler",
|
||||
InputSchema = JsonSerializer.Deserialize<JsonElement>("""
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": []
|
||||
}
|
||||
"""),
|
||||
}],
|
||||
};
|
||||
|
||||
default:
|
||||
|
|
@ -117,43 +109,13 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
|||
})
|
||||
.WithTools<EchoTool>();
|
||||
|
||||
sc.AddSingleton(new ObjectWithId());
|
||||
_serviceProvider = sc.BuildServiceProvider();
|
||||
|
||||
var server = _serviceProvider.GetRequiredService<IMcpServer>();
|
||||
_cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken);
|
||||
_serverTask = server.RunAsync(cancellationToken: _cts.Token);
|
||||
}
|
||||
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
await _cts.CancelAsync();
|
||||
|
||||
_clientToServerPipe.Writer.Complete();
|
||||
_serverToClientPipe.Writer.Complete();
|
||||
|
||||
await _serverTask;
|
||||
|
||||
await _serviceProvider.DisposeAsync();
|
||||
_cts.Dispose();
|
||||
Dispose();
|
||||
}
|
||||
|
||||
private async Task<IMcpClient> CreateMcpClientForServer(McpClientOptions? options = null)
|
||||
{
|
||||
return await McpClientFactory.CreateAsync(
|
||||
new StreamClientTransport(
|
||||
serverInput: _clientToServerPipe.Writer.AsStream(),
|
||||
_serverToClientPipe.Reader.AsStream(),
|
||||
LoggerFactory),
|
||||
loggerFactory: LoggerFactory,
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
services.AddSingleton(new ObjectWithId());
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public void Adds_Tools_To_Server()
|
||||
{
|
||||
var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var serverOptions = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var tools = serverOptions.Capabilities?.Tools?.ToolCollection;
|
||||
Assert.NotNull(tools);
|
||||
Assert.NotEmpty(tools);
|
||||
|
|
@ -183,8 +145,8 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
|||
[Fact]
|
||||
public async Task Can_Create_Multiple_Servers_From_Options_And_List_Registered_Tools()
|
||||
{
|
||||
var options = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var loggerFactory = _serviceProvider.GetRequiredService<ILoggerFactory>();
|
||||
var options = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var loggerFactory = ServiceProvider.GetRequiredService<ILoggerFactory>();
|
||||
|
||||
for (int i = 0; i < 2; i++)
|
||||
{
|
||||
|
|
@ -192,7 +154,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
|||
var stdoutPipe = new Pipe();
|
||||
|
||||
await using var transport = new StreamServerTransport(stdinPipe.Reader.AsStream(), stdoutPipe.Writer.AsStream());
|
||||
await using var server = McpServerFactory.Create(transport, options, loggerFactory, _serviceProvider);
|
||||
await using var server = McpServerFactory.Create(transport, options, loggerFactory, ServiceProvider);
|
||||
var serverRunTask = server.RunAsync(TestContext.Current.CancellationToken);
|
||||
|
||||
await using (var client = await McpClientFactory.CreateAsync(
|
||||
|
|
@ -237,7 +199,7 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
|||
var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
|
||||
Assert.False(notificationRead.IsCompleted);
|
||||
|
||||
var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var serverOptions = ServiceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
|
||||
var serverTools = serverOptions.Capabilities?.Tools?.ToolCollection;
|
||||
Assert.NotNull(serverTools);
|
||||
|
||||
|
|
@ -444,7 +406,9 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
|||
[Fact]
|
||||
public void WithTools_InvalidArgs_Throws()
|
||||
{
|
||||
Assert.Throws<ArgumentNullException>("toolTypes", () => _builder.WithTools((IEnumerable<Type>)null!));
|
||||
IMcpServerBuilder builder = new ServiceCollection().AddMcpServer();
|
||||
|
||||
Assert.Throws<ArgumentNullException>("toolTypes", () => builder.WithTools((IEnumerable<Type>)null!));
|
||||
|
||||
IMcpServerBuilder nullBuilder = null!;
|
||||
Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithTools<object>());
|
||||
|
|
@ -455,9 +419,11 @@ public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
|
|||
[Fact]
|
||||
public void Empty_Enumerables_Is_Allowed()
|
||||
{
|
||||
_builder.WithTools(toolTypes: []); // no exception
|
||||
_builder.WithTools<object>(); // no exception even though no tools exposed
|
||||
_builder.WithToolsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no tools exposed
|
||||
IMcpServerBuilder builder = new ServiceCollection().AddMcpServer();
|
||||
|
||||
builder.WithTools(toolTypes: []); // no exception
|
||||
builder.WithTools<object>(); // no exception even though no tools exposed
|
||||
builder.WithToolsFromAssembly(typeof(AIFunction).Assembly); // no exception even though no tools exposed
|
||||
}
|
||||
|
||||
[Fact]
|
||||
|
|
|
|||
Loading…
Reference in a new issue