Overhaul tool handling (#89)

* Overhaul tool handling

- Renames [McpTool{Type}] to [McpServerTool{Type}], in order to distinguish representations of tools on the server from tools on the client.
- Enables [McpServerTool] methods to have arguments injected from DI, as well as a specific set of types from the implementation directly, like IMcpServer.
- Renames WithTools to WithToolsFromAssembly.
- All of the WithToolsXx methods now publish each individual McpServerTool into DI; other code can do so as well. The options setup code gathers all of the tools from DI and combines them into a collection which is then stored in McpServerOptions and used to construct the server.
- The server tools specified via DI as well as manually-provided handlers, using CallToolHandler as a fallback in case the requested tool doesn't exist in the tools collection, and ListToolsHandler augments the information for whatever tools exist.
- The tools are stored in McpServerOptions in a new McpServerToolCollection type, which is a thread-safe container that exposes add/removal notification. Adding/removing tools triggers a change notification that in turn sends a notification to the client about a tools update.
- The ServerOptions are exposed from the server instance.
- Removed cursor-based APIs from McpClientExtensions.
- Changed McpClientExtensions APIs to return `Task<IList<...>` rather than `IAsyncEnumerable<...>`.

I'm sure this will need to evolve further before we're "done", but it's a significant improvement from where we are now. One area I'm not super happy with is the notifying collection; that might work ok for a stdio server, where you can grab the created collection from the server's options and mutate it, but I don't know how well that's going to work for sse servers.

* Make McpServerTool{Collection} extensible

* Fix README.md example

* Add notifications test

* Address feedback
This commit is contained in:
Stephen Toub 2025-03-25 15:32:46 -04:00 • committed by GitHub
commit caca2ae61d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 2703 additions and 686 deletions

View file

@ -0,0 +1,86 @@
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.Client;
using ModelContextProtocol.Configuration;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Transport;
using System.IO.Pipelines;
namespace ModelContextProtocol.Tests.Client;
public class McpClientExtensionsTests
{
private Pipe _clientToServerPipe = new();
private Pipe _serverToClientPipe = new();
private readonly IMcpServer _server;
public McpClientExtensionsTests()
{
ServiceCollection sc = new();
sc.AddSingleton<IServerTransport>(new StdioServerTransport("TestServer", _clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream()));
sc.AddMcpServer();
for (int f = 0; f < 10; f++)
{
string name = $"Method{f}";
sc.AddSingleton(McpServerTool.Create((int i) => $"{name} Result {i}", name));
}
_server = sc.BuildServiceProvider().GetRequiredService<IMcpServer>();
}
public ValueTask DisposeAsync()
{
_clientToServerPipe.Writer.Complete();
_serverToClientPipe.Writer.Complete();
return _server.DisposeAsync();
}
private async Task<IMcpClient> CreateMcpClientForServer()
{
await _server.StartAsync(TestContext.Current.CancellationToken);
var stdin = new StreamReader(_serverToClientPipe.Reader.AsStream());
var stdout = new StreamWriter(_clientToServerPipe.Writer.AsStream());
var serverConfig = new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
};
return await McpClientFactory.CreateAsync(
serverConfig,
createTransportFunc: (_, _) => new StreamClientTransport(stdin, stdout),
cancellationToken: TestContext.Current.CancellationToken);
}
[Fact]
public async Task ListToolsAsync_AllToolsReturned()
{
IMcpClient client = await CreateMcpClientForServer();
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
Assert.Equal(10, tools.Count);
var echo = tools.Single(t => t.Name == "Method4");
var result = await echo.InvokeAsync(new Dictionary<string, object?>() { ["i"] = 42 }, TestContext.Current.CancellationToken);
Assert.Contains("Method4 Result 42", result?.ToString());
}
[Fact]
public async Task EnumerateToolsAsync_AllToolsReturned()
{
IMcpClient client = await CreateMcpClientForServer();
await foreach (var tool in client.EnumerateToolsAsync(TestContext.Current.CancellationToken))
{
if (tool.Name == "Method4")
{
var result = await tool.InvokeAsync(new Dictionary<string, object?>() { ["i"] = 42 }, TestContext.Current.CancellationToken);
Assert.Contains("Method4 Result 42", result?.ToString());
return;
}
}
Assert.Fail("Couldn't find target method");
}
}

View file

@ -7,7 +7,6 @@ using System.Text.Json;
using ModelContextProtocol.Configuration;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Tests.Utils;
using Xunit.Sdk;
using System.Text.Encodings.Web;
using System.Text.Json.Serialization.Metadata;
using System.Text.Json.Serialization;
@ -18,6 +17,8 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
{
private static readonly string? s_openAIKey = Environment.GetEnvironmentVariable("AI:OpenAI:ApiKey")!;
public static bool NoOpenAIKeySet => string.IsNullOrWhiteSpace(s_openAIKey);
private readonly ClientIntegrationTestFixture _fixture;
public ClientIntegrationTests(ClientIntegrationTestFixture fixture, ITestOutputHelper testOutputHelper)
@ -68,13 +69,10 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
// act
await using var client = await _fixture.CreateClientAsync(clientId);
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken).ToListAsync(TestContext.Current.CancellationToken);
var aiFunctions = await client.GetAIFunctionsAsync(TestContext.Current.CancellationToken);
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
// assert
Assert.NotEmpty(tools);
Assert.NotEmpty(aiFunctions);
Assert.Equal(tools.Count, aiFunctions.Count);
}
[Theory]
@ -87,7 +85,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
await using var client = await _fixture.CreateClientAsync(clientId);
var result = await client.CallToolAsync(
"echo",
new Dictionary<string, object>
new Dictionary<string, object?>
{
["message"] = "Hello MCP!"
},
@ -109,7 +107,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
// act
await using var client = await _fixture.CreateClientAsync(clientId);
var aiFunctions = await client.GetAIFunctionsAsync(TestContext.Current.CancellationToken);
var aiFunctions = await client.ListToolsAsync(TestContext.Current.CancellationToken);
var echo = aiFunctions.Single(t => t.Name == "echo");
var result = await echo.InvokeAsync([new KeyValuePair<string, object?>("message", "Hello MCP!")], TestContext.Current.CancellationToken);
@ -126,7 +124,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
// act
await using var client = await _fixture.CreateClientAsync(clientId);
var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken).ToListAsync(TestContext.Current.CancellationToken);
var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
// assert
Assert.NotEmpty(prompts);
@ -158,7 +156,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
// act
await using var client = await _fixture.CreateClientAsync(clientId);
var arguments = new Dictionary<string, object>
var arguments = new Dictionary<string, object?>
{
{ "temperature", "0.7" },
{ "style", "formal" }
@ -191,7 +189,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
// act
await using var client = await _fixture.CreateClientAsync(clientId);
List<ResourceTemplate> allResourceTemplates = await client.ListResourceTemplatesAsync(TestContext.Current.CancellationToken).ToListAsync(TestContext.Current.CancellationToken);
IList<ResourceTemplate> allResourceTemplates = await client.ListResourceTemplatesAsync(TestContext.Current.CancellationToken);
// The server provides a single test resource template
Assert.Single(allResourceTemplates);
@ -206,15 +204,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
// act
await using var client = await _fixture.CreateClientAsync(clientId);
List<Resource> allResources = [];
string? cursor = null;
do
{
var resources = await client.ListResourcesAsync(cursor, CancellationToken.None);
allResources.AddRange(resources.Resources);
cursor = resources.NextCursor;
}
while (cursor != null);
IList<Resource> allResources = await client.ListResourcesAsync(TestContext.Current.CancellationToken);
// The server provides 100 test resources
Assert.Equal(100, allResources.Count);
@ -393,7 +383,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
// Call the server's sampleLLM tool which should trigger our sampling handler
var result = await client.CallToolAsync(
"sampleLLM",
new Dictionary<string, object>
new Dictionary<string, object?>
{
["prompt"] = "Test prompt",
["maxTokens"] = 100
@ -483,7 +473,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
// act
var result = await client.CallToolAsync(
"read_graph",
[],
new Dictionary<string, object?>(),
TestContext.Current.CancellationToken);
// assert
@ -494,17 +484,15 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
await client.DisposeAsync();
}
[Fact]
public async Task GetAIFunctionsAsync_UsingEverythingServer_ToolsAreProperlyCalled()
[Fact(Skip = "Requires OpenAI API Key", SkipWhen = nameof(NoOpenAIKeySet))]
public async Task ListToolsAsync_UsingEverythingServer_ToolsAreProperlyCalled()
{
SkipTestIfNoOpenAIKey();
// Get the MCP client and tools from it.
await using var client = await McpClientFactory.CreateAsync(
_fixture.EverythingServerConfig,
_fixture.DefaultOptions,
cancellationToken: TestContext.Current.CancellationToken);
var mappedTools = await client.GetAIFunctionsAsync(TestContext.Current.CancellationToken);
var mappedTools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
// Create the chat client.
using IChatClient chatClient = new OpenAIClient(s_openAIKey).AsChatClient("gpt-4o-mini")
@ -527,11 +515,9 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
Assert.Contains("Echo: Hello MCP!", response.Text);
}
[Fact]
[Fact(Skip = "Requires OpenAI API Key", SkipWhen = nameof(NoOpenAIKeySet))]
public async Task SamplingViaChatClient_RequestResponseProperlyPropagated()
{
SkipTestIfNoOpenAIKey();
await using var client = await McpClientFactory.CreateAsync(_fixture.EverythingServerConfig, new()
{
ClientInfo = new() { Name = nameof(SamplingViaChatClient_RequestResponseProperlyPropagated), Version = "1.0.0" },
@ -544,7 +530,7 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
},
}, cancellationToken: TestContext.Current.CancellationToken);
var result = await client.CallToolAsync("sampleLLM", new()
var result = await client.CallToolAsync("sampleLLM", new Dictionary<string, object?>()
{
["prompt"] = "In just a few words, what is the most famous tower in Paris?",
}, TestContext.Current.CancellationToken);
@ -590,9 +576,4 @@ public class ClientIntegrationTests : LoggedTest, IClassFixture<ClientIntegratio
// assert
Assert.True(logCounter > 0);
}
private static void SkipTestIfNoOpenAIKey()
{
Assert.SkipWhen(s_openAIKey is null, "No OpenAI key provided. Skipping test.");
}
}

View file

@ -1,72 +1,139 @@
using System.ComponentModel;
using System.Text.Json;
using System.Text.RegularExpressions;
using ModelContextProtocol.Configuration;
using ModelContextProtocol.Server;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Options;
using Moq;
using ModelContextProtocol.Protocol.Transport;
using System.IO.Pipelines;
using ModelContextProtocol.Client;
using ModelContextProtocol.Configuration;
using ModelContextProtocol.Tests.Transport;
using System.Text.RegularExpressions;
using Microsoft.Extensions.AI;
using System.Threading.Channels;
using ModelContextProtocol.Protocol.Messages;
namespace ModelContextProtocol.Tests.Configuration;
public class McpServerBuilderExtensionsToolsTests
public class McpServerBuilderExtensionsToolsTests : IAsyncDisposable
{
private readonly Mock<IMcpServerBuilder> _builder;
private readonly ServiceCollection _services;
private Pipe _clientToServerPipe = new();
private Pipe _serverToClientPipe = new();
private readonly IMcpServerBuilder _builder;
private readonly IMcpServer _server;
public McpServerBuilderExtensionsToolsTests()
{
_services = new ServiceCollection();
_builder = new Mock<IMcpServerBuilder>();
_builder.SetupGet(b => b.Services).Returns(_services);
ServiceCollection sc = new();
sc.AddSingleton<IServerTransport>(new StdioServerTransport("TestServer", _clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream()));
_builder = sc.AddMcpServer().WithTools<EchoTool>();
_server = sc.BuildServiceProvider().GetRequiredService<IMcpServer>();
}
public ValueTask DisposeAsync()
{
_clientToServerPipe.Writer.Complete();
_serverToClientPipe.Writer.Complete();
return _server.DisposeAsync();
}
private async Task<IMcpClient> CreateMcpClientForServer()
{
await _server.StartAsync(TestContext.Current.CancellationToken);
var stdin = new StreamReader(_serverToClientPipe.Reader.AsStream());
var stdout = new StreamWriter(_clientToServerPipe.Writer.AsStream());
var serverConfig = new McpServerConfig()
{
Id = "TestServer",
Name = "TestServer",
TransportType = "ignored",
};
return await McpClientFactory.CreateAsync(
serverConfig,
createTransportFunc: (_, _) => new StreamClientTransport(stdin, stdout),
cancellationToken: TestContext.Current.CancellationToken);
}
[Fact]
public void Adds_Tools_To_Server()
{
_builder.Object.WithTools(typeof(EchoTool));
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
Assert.NotNull(options.ListToolsHandler);
Assert.NotNull(options.CallToolHandler);
var tools = _server.ServerOptions?.Capabilities?.Tools?.ToolCollection;
Assert.NotNull(tools);
Assert.NotEmpty(tools);
}
[Fact]
public async Task Can_List_Registered_Tool()
public async Task Can_List_Registered_Tools()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
Assert.Equal(10, tools.Count);
var result = await options.ListToolsHandler!(new(Mock.Of<IMcpServer>(), new()), CancellationToken.None);
Assert.NotNull(result);
Assert.NotEmpty(result.Tools);
McpClientTool echoTool = tools.First(t => t.Name == "Echo");
Assert.Equal("Echo", echoTool.Name);
Assert.Equal("Echoes the input back to the client.", echoTool.Description);
Assert.Equal("object", echoTool.JsonSchema.GetProperty("type").GetString());
Assert.Equal(JsonValueKind.Object, echoTool.JsonSchema.GetProperty("properties").GetProperty("message").ValueKind);
Assert.Equal("the echoes message", echoTool.JsonSchema.GetProperty("properties").GetProperty("message").GetProperty("description").GetString());
Assert.Equal(1, echoTool.JsonSchema.GetProperty("required").GetArrayLength());
var tool = result.Tools[0];
Assert.Equal("Echo", tool.Name);
Assert.Equal("Echoes the input back to the client.", tool.Description);
Assert.Equal("object", tool.InputSchema.GetProperty("type").GetString());
Assert.Equal(JsonValueKind.Object, tool.InputSchema.GetProperty("properties").GetProperty("message").ValueKind);
Assert.Equal("the echoes message", tool.InputSchema.GetProperty("properties").GetProperty("message").GetProperty("description").GetString());
Assert.Equal(1, tool.InputSchema.GetProperty("required").GetArrayLength());
McpClientTool doubleEchoTool = tools.First(t => t.Name == "double_echo");
Assert.Equal("double_echo", doubleEchoTool.Name);
Assert.Equal("Echoes the input back to the client.", doubleEchoTool.Description);
}
tool = result.Tools[1];
Assert.Equal("double_echo", tool.Name);
Assert.Equal("Echoes the input back to the client.", tool.Description);
[Fact]
public async Task Can_Be_Notified_Of_Tool_Changes()
{
IMcpClient client = await CreateMcpClientForServer();
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
Assert.Equal(10, tools.Count);
Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
client.AddNotificationHandler("notifications/tools/list_changed", notification =>
{
listChanged.Writer.TryWrite(notification);
return Task.CompletedTask;
});
var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
var serverTools = _server.ServerOptions.Capabilities?.Tools?.ToolCollection;
Assert.NotNull(serverTools);
var newTool = McpServerTool.Create([McpServerTool(name: "NewTool")] () => "42");
serverTools.Add(newTool);
await notificationRead;
tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
Assert.Equal(11, tools.Count);
Assert.Contains(tools, t => t.Name == "NewTool");
notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
Assert.False(notificationRead.IsCompleted);
serverTools.Remove(newTool);
await notificationRead;
tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
Assert.Equal(10, tools.Count);
Assert.DoesNotContain(tools, t => t.Name == "NewTool");
}
[Fact]
public async Task Can_Call_Registered_Tool()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var result = await client.CallToolAsync(
"Echo",
new Dictionary<string, object?>() { ["message"] = "Peter" },
TestContext.Current.CancellationToken);
var result = await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "Echo", Arguments = new() { { "message", "Peter" } } }), CancellationToken.None);
Assert.NotNull(result);
Assert.NotNull(result.Content);
Assert.NotEmpty(result.Content);
@ -78,13 +145,13 @@ public class McpServerBuilderExtensionsToolsTests
[Fact]
public async Task Can_Call_Registered_Tool_With_Array_Result()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var result = await client.CallToolAsync(
"EchoArray",
new Dictionary<string, object?>() { ["message"] = "Peter" },
TestContext.Current.CancellationToken);
var result = await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "EchoArray", Arguments = new() { { "message", "Peter" } } }), CancellationToken.None);
Assert.NotNull(result);
Assert.NotNull(result.Content);
Assert.NotEmpty(result.Content);
@ -95,12 +162,12 @@ public class McpServerBuilderExtensionsToolsTests
[Fact]
public async Task Can_Call_Registered_Tool_With_Null_Result()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var result = await client.CallToolAsync(
"ReturnNull",
cancellationToken: TestContext.Current.CancellationToken);
var result = await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "ReturnNull" }), CancellationToken.None);
Assert.NotNull(result);
Assert.NotNull(result.Content);
Assert.Empty(result.Content);
@ -109,30 +176,29 @@ public class McpServerBuilderExtensionsToolsTests
[Fact]
public async Task Can_Call_Registered_Tool_With_Json_Result()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var result = await client.CallToolAsync(
"ReturnJson",
cancellationToken: TestContext.Current.CancellationToken);
var result = await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "ReturnJson" }), CancellationToken.None);
Assert.NotNull(result);
Assert.NotNull(result.Content);
Assert.NotEmpty(result.Content);
Assert.Equal("{\"SomeProp\":false}", Regex.Replace(result.Content[0].Text ?? string.Empty, "\\s+", ""));
Assert.Equal("""{"SomeProp":false}""", Regex.Replace(result.Content[0].Text ?? string.Empty, "\\s+", ""));
Assert.Equal("text", result.Content[0].Type);
}
[Fact]
public async Task Can_Call_Registered_Tool_With_Int_Result()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var result = await client.CallToolAsync(
"ReturnInteger",
cancellationToken: TestContext.Current.CancellationToken);
var result = await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "ReturnInteger" }), CancellationToken.None);
Assert.NotNull(result);
Assert.NotNull(result.Content);
Assert.NotEmpty(result.Content);
@ -140,51 +206,16 @@ public class McpServerBuilderExtensionsToolsTests
Assert.Equal("text", result.Content[0].Type);
}
[Fact]
public async Task Can_Call_Registered_Tool_And_Pass_Cancellation_Token()
{
_builder.Object.WithTools(typeof(EchoTool));
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
using var cts = new CancellationTokenSource();
var token = cts.Token;
var result = await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "ReturnCancellationToken" }), token);
Assert.NotNull(result);
Assert.NotNull(result.Content);
Assert.NotEmpty(result.Content);
Assert.Equal(token.GetHashCode().ToString(), result.Content[0].Text);
}
[Fact]
public async Task Can_Call_Registered_Tool_And_Returns_Cancelled_Response()
{
_builder.Object.WithTools(typeof(EchoTool));
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
using var cts = new CancellationTokenSource();
var token = cts.Token;
await cts.CancelAsync();
var action = async () => await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "ReturnCancellationToken" }), token);
await Assert.ThrowsAsync<OperationCanceledException>(action);
}
[Fact]
public async Task Can_Call_Registered_Tool_And_Pass_ComplexType()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var result = await client.CallToolAsync(
"EchoComplex",
new Dictionary<string, object?>() { ["complex"] = JsonDocument.Parse("""{"Name": "Peter", "Age": 25}""").RootElement },
cancellationToken: TestContext.Current.CancellationToken);
var result = await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "EchoComplex", Arguments = new() { { "complex", JsonDocument.Parse("{\"Name\": \"Peter\", \"Age\": 25}").RootElement } } }), CancellationToken.None);
Assert.NotNull(result);
Assert.NotNull(result.Content);
Assert.NotEmpty(result.Content);
@ -196,166 +227,195 @@ public class McpServerBuilderExtensionsToolsTests
[Fact]
public async Task Returns_IsError_Content_When_Tool_Fails()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var result = await client.CallToolAsync(
"ThrowException",
cancellationToken: TestContext.Current.CancellationToken);
var response = await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = nameof(EchoTool.ThrowException) }), CancellationToken.None);
Assert.True(response.IsError);
Assert.NotNull(response.Content);
Assert.NotEmpty(response.Content);
Assert.Contains("Test error", response.Content[0].Text);
Assert.True(result.IsError);
Assert.NotNull(result.Content);
Assert.NotEmpty(result.Content);
Assert.Contains("Test error", result.Content[0].Text);
}
[Fact]
public async Task Throws_Exception_On_Unknown_Tool()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var e = await Assert.ThrowsAsync<McpClientException>(async () => await client.CallToolAsync(
"NotRegisteredTool",
cancellationToken: TestContext.Current.CancellationToken));
var exception = await Assert.ThrowsAsync<McpServerException>(async () => await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "NotRegisteredTool" }), CancellationToken.None));
Assert.Contains("'NotRegisteredTool'", exception.Message);
Assert.Contains("'NotRegisteredTool'", e.Message);
}
[Fact(Skip = "https://github.com/dotnet/extensions/issues/6124")]
public async Task Throws_Exception_Missing_Parameter()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var e = await Assert.ThrowsAsync<McpClientException>(async () => await client.CallToolAsync(
"Echo",
cancellationToken: TestContext.Current.CancellationToken));
var exception = await Assert.ThrowsAsync<McpServerException>(async () => await options.CallToolHandler!(new(Mock.Of<IMcpServer>(), new() { Name = "Echo" }), CancellationToken.None));
Assert.Equal("Missing required argument 'message'.", exception.Message);
Assert.Equal("Missing required argument 'message'.", e.Message);
}
[Fact]
public void Throws_Exception_For_Null_Types()
public void WithTools_InvalidArgs_Throws()
{
Assert.Throws<ArgumentNullException>("toolTypes", () => _builder.Object.WithTools(toolTypes: null!));
Assert.Throws<ArgumentNullException>("toolTypes", () => _builder.WithTools((IEnumerable<Type>)null!));
IMcpServerBuilder nullBuilder = null!;
Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithTools<object>());
Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithTools(Array.Empty<Type>()));
Assert.Throws<ArgumentNullException>("builder", () => nullBuilder.WithToolsFromAssembly());
}
[Fact]
public void Empty_Types_Is_Allowed()
public void Empty_Enumerables_Is_Allowed()
{
_builder.Object.WithTools(toolTypes: []); // no exception
_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]
public async Task Register_Tools_From_Current_Assembly()
public void Register_Tools_From_Current_Assembly()
{
_builder.Object.WithTools();
ServiceCollection sc = new();
sc.AddMcpServer().WithToolsFromAssembly();
IServiceProvider services = sc.BuildServiceProvider();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var result = await options.ListToolsHandler!(new(Mock.Of<IMcpServer>(), new()), CancellationToken.None);
Assert.NotNull(result);
Assert.NotEmpty(result.Tools);
var tool = result.Tools[0];
Assert.Equal("Echo", tool.Name);
}
[Fact]
public void Ok_If_No_Tools_Are_Found_In_Given_Assembly()
{
_builder.Object.WithToolsFromAssembly(typeof(Mock).Assembly);
Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "Echo");
}
[Fact]
public async Task Recognizes_Parameter_Types()
{
_builder.Object.WithTools(typeof(EchoTool));
IMcpClient client = await CreateMcpClientForServer();
var serviceProvider = _services.BuildServiceProvider();
var options = serviceProvider.GetRequiredService<IOptions<McpServerHandlers>>().Value;
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
var result = await options.ListToolsHandler!(new(Mock.Of<IMcpServer>(), new()), CancellationToken.None);
Assert.NotNull(result);
Assert.NotEmpty(result.Tools);
Assert.NotNull(tools);
Assert.NotEmpty(tools);
var tool = result.Tools.First(t => t.Name == "TestTool");
var tool = tools.First(t => t.Name == "TestTool");
Assert.Equal("TestTool", tool.Name);
Assert.Empty(tool.Description!);
Assert.Equal("object", tool.InputSchema.GetProperty("type").GetString());
Assert.Equal("object", tool.JsonSchema.GetProperty("type").GetString());
Assert.Contains("integer", tool.InputSchema.GetProperty("properties").GetProperty("number").GetProperty("type").GetString());
Assert.Contains("number", tool.InputSchema.GetProperty("properties").GetProperty("otherNumber").GetProperty("type").GetString());
Assert.Contains("boolean", tool.InputSchema.GetProperty("properties").GetProperty("someCheck").GetProperty("type").GetString());
Assert.Contains("string", tool.InputSchema.GetProperty("properties").GetProperty("someDate").GetProperty("type").GetString());
Assert.Contains("string", tool.InputSchema.GetProperty("properties").GetProperty("someOtherDate").GetProperty("type").GetString());
Assert.Contains("array", tool.InputSchema.GetProperty("properties").GetProperty("data").GetProperty("type").GetString());
Assert.Contains("object", tool.InputSchema.GetProperty("properties").GetProperty("complexObject").GetProperty("type").GetString());
Assert.Contains("integer", tool.JsonSchema.GetProperty("properties").GetProperty("number").GetProperty("type").GetString());
Assert.Contains("number", tool.JsonSchema.GetProperty("properties").GetProperty("otherNumber").GetProperty("type").GetString());
Assert.Contains("boolean", tool.JsonSchema.GetProperty("properties").GetProperty("someCheck").GetProperty("type").GetString());
Assert.Contains("string", tool.JsonSchema.GetProperty("properties").GetProperty("someDate").GetProperty("type").GetString());
Assert.Contains("string", tool.JsonSchema.GetProperty("properties").GetProperty("someOtherDate").GetProperty("type").GetString());
Assert.Contains("array", tool.JsonSchema.GetProperty("properties").GetProperty("data").GetProperty("type").GetString());
Assert.Contains("object", tool.JsonSchema.GetProperty("properties").GetProperty("complexObject").GetProperty("type").GetString());
}
[McpToolType]
public static class EchoTool
[Fact]
public void Register_Tools_From_Multiple_Sources()
{
[McpTool, Description("Echoes the input back to the client.")]
ServiceCollection sc = new();
sc.AddMcpServer()
.WithTools<EchoTool>()
.WithTools<AnotherToolType>()
.WithTools(typeof(ToolTypeWithNoAttribute));
IServiceProvider services = sc.BuildServiceProvider();
Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "double_echo");
Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "DifferentName");
Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "MethodB");
Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "MethodC");
Assert.Contains(services.GetServices<McpServerTool>(), t => t.ProtocolTool.Name == "MethodD");
}
[McpServerToolType]
public sealed class EchoTool
{
[McpServerTool, Description("Echoes the input back to the client.")]
public static string Echo([Description("the echoes message")] string message)
{
return "hello " + message;
}
[McpTool("double_echo"), Description("Echoes the input back to the client.")]
[McpServerTool("double_echo"), Description("Echoes the input back to the client.")]
public static string Echo2(string message)
{
return "hello hello" + message;
}
[McpTool]
[McpServerTool]
public static string TestTool(int number, double otherNumber, bool someCheck, DateTime someDate, DateTimeOffset someOtherDate, string[] data, ComplexObject complexObject)
{
return "hello hello";
}
[McpTool]
[McpServerTool]
public static string[] EchoArray(string message)
{
return ["hello " + message, "hello2 " + message];
}
[McpTool]
[McpServerTool]
public static string? ReturnNull()
{
return null;
}
[McpTool]
[McpServerTool]
public static JsonElement ReturnJson()
{
return JsonDocument.Parse("{\"SomeProp\": false}").RootElement;
}
[McpTool]
[McpServerTool]
public static int ReturnInteger()
{
return 5;
}
[McpTool]
[McpServerTool]
public static string ThrowException()
{
throw new InvalidOperationException("Test error");
}
[McpTool]
[McpServerTool]
public static int ReturnCancellationToken(CancellationToken cancellationToken)
{
return cancellationToken.GetHashCode();
}
[McpTool]
[McpServerTool]
public static string EchoComplex(ComplexObject complex)
{
return complex.Name!;
}
}
[McpServerToolType]
internal class AnotherToolType
{
[McpServerTool("DifferentName")]
private static string MethodA(int a) => a.ToString();
[McpServerTool]
internal static string MethodB(string b) => b.ToString();
[McpServerTool]
protected static string MethodC(long c) => c.ToString();
}
internal class ToolTypeWithNoAttribute
{
[McpServerTool]
public static string MethodD(string d) => d.ToString();
}
public class ComplexObject
{
public string? Name { get; set; }

View file

@ -1,51 +1,51 @@
<Project Sdk="Microsoft.NET.Sdk">
<PropertyGroup>
<TargetFramework>net8.0</TargetFramework>
<ImplicitUsings>enable</ImplicitUsings>
<Nullable>enable</Nullable>
<LangVersion>Latest</LangVersion>
<IsPackable>false</IsPackable>
<IsTestProject>true</IsTestProject>
<RootNamespace>ModelContextProtocol.Tests</RootNamespace>
</PropertyGroup>
<ItemGroup>
<Project Sdk="Microsoft.NET.Sdk">
<PropertyGroup>
<TargetFramework>net8.0</TargetFramework>
<ImplicitUsings>enable</ImplicitUsings>
<Nullable>enable</Nullable>
<LangVersion>Latest</LangVersion>
<IsPackable>false</IsPackable>
<IsTestProject>true</IsTestProject>
<RootNamespace>ModelContextProtocol.Tests</RootNamespace>
</PropertyGroup>
<ItemGroup>
<PackageReference Include="coverlet.collector">
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
<PrivateAssets>all</PrivateAssets>
</PackageReference>
</PackageReference>
<PackageReference Include="GitHubActionsTestLogger">
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
<PrivateAssets>all</PrivateAssets>
</PackageReference>
<PackageReference Include="Microsoft.Extensions.AI.OpenAI" />
<PackageReference Include="Microsoft.Extensions.Logging" />
<PackageReference Include="Microsoft.Extensions.Logging.Console" />
<PackageReference Include="Microsoft.NET.Test.Sdk" />
<PackageReference Include="Moq" />
<PackageReference Include="System.Linq.AsyncEnumerable" />
<PackageReference Include="xunit.v3" />
<PackageReference Include="xunit.runner.visualstudio">
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
<PrivateAssets>all</PrivateAssets>
</PackageReference>
</ItemGroup>
<ItemGroup>
<ProjectReference Include="..\..\src\ModelContextProtocol\ModelContextProtocol.csproj" />
<ProjectReference Include="..\ModelContextProtocol.TestServer\ModelContextProtocol.TestServer.csproj" />
<ProjectReference Include="..\ModelContextProtocol.TestSseServer\ModelContextProtocol.TestSseServer.csproj" />
</ItemGroup>
<ItemGroup>
<Content Condition="$([MSBuild]::IsOSPlatform('Windows'))" Include="..\ModelContextProtocol.TestServer\bin\$(Configuration)\net8.0\TestServer.exe">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Condition="$([MSBuild]::IsOSPlatform('Linux'))" Include="..\ModelContextProtocol.TestServer\bin\$(Configuration)\net8.0\TestServer.dll">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
</ItemGroup>
</Project>
</PackageReference>
<PackageReference Include="Microsoft.Extensions.AI.OpenAI" />
<PackageReference Include="Microsoft.Extensions.Logging" />
<PackageReference Include="Microsoft.Extensions.Logging.Console" />
<PackageReference Include="Microsoft.NET.Test.Sdk" />
<PackageReference Include="Moq" />
<PackageReference Include="System.Linq.AsyncEnumerable" />
<PackageReference Include="xunit.v3" />
<PackageReference Include="xunit.runner.visualstudio">
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
<PrivateAssets>all</PrivateAssets>
</PackageReference>
</ItemGroup>
<ItemGroup>
<ProjectReference Include="..\..\src\ModelContextProtocol\ModelContextProtocol.csproj" />
<ProjectReference Include="..\ModelContextProtocol.TestServer\ModelContextProtocol.TestServer.csproj" />
<ProjectReference Include="..\ModelContextProtocol.TestSseServer\ModelContextProtocol.TestSseServer.csproj" />
</ItemGroup>
<ItemGroup>
<Content Condition="$([MSBuild]::IsOSPlatform('Windows'))" Include="..\ModelContextProtocol.TestServer\bin\$(Configuration)\net8.0\TestServer.exe">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
<Content Condition="$([MSBuild]::IsOSPlatform('Linux'))" Include="..\ModelContextProtocol.TestServer\bin\$(Configuration)\net8.0\TestServer.dll">
<CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>
</Content>
</ItemGroup>
</Project>

View file

@ -1,5 +1,4 @@
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Utils.Json;
using System.Text.Json;
namespace ModelContextProtocol.Tests.Protocol;

View file

@ -1,4 +1,5 @@
using Microsoft.Extensions.AI;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol.Messages;
@ -24,7 +25,7 @@ public class McpServerTests : LoggedTest
_serverTransport = new Mock<IServerTransport>();
_logger = new Mock<ILogger>();
_options = CreateOptions();
_serviceProvider = new Mock<IServiceProvider>().Object;
_serviceProvider = new ServiceCollection().BuildServiceProvider();
}
private static McpServerOptions CreateOptions(ServerCapabilities? capabilities = null)
@ -133,8 +134,7 @@ public class McpServerTests : LoggedTest
await transport.SendMessageAsync(new JsonRpcNotification
{
Method = "notifications/initialized"
}
, TestContext.Current.CancellationToken);
}, TestContext.Current.CancellationToken);
await Task.Delay(50, TestContext.Current.CancellationToken);
@ -678,12 +678,15 @@ public class McpServerTests : LoggedTest
public bool IsInitialized => throw new NotImplementedException();
public Implementation? ClientInfo => throw new NotImplementedException();
public IServiceProvider? ServiceProvider => throw new NotImplementedException();
public McpServerOptions ServerOptions => throw new NotImplementedException();
public IServiceProvider? Services => throw new NotImplementedException();
public void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler) =>
throw new NotImplementedException();
public Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default) =>
throw new NotImplementedException();
public Task StartAsync(CancellationToken cancellationToken = default) =>
throw new NotImplementedException();
public object? GetService(Type serviceType, object? serviceKey = null) => null;
}
}

View file

@ -0,0 +1,93 @@
using Microsoft.Extensions.DependencyInjection;
using ModelContextProtocol.Protocol.Types;
using ModelContextProtocol.Server;
using Moq;
using System.Reflection;
using System.Text.Json;
namespace ModelContextProtocol.Tests.Server;
public class McpServerToolTests
{
[Fact]
public void Create_InvalidArgs_Throws()
{
Assert.Throws<ArgumentNullException>("function", () => McpServerTool.Create(null!));
Assert.Throws<ArgumentNullException>("method", () => McpServerTool.Create((MethodInfo)null!));
Assert.Throws<ArgumentNullException>("method", () => McpServerTool.Create((Delegate)null!));
}
[Fact]
public async Task SupportsIMcpServer()
{
Mock<IMcpServer> mockServer = new();
McpServerTool tool = McpServerTool.Create((IMcpServer server) =>
{
Assert.Same(mockServer.Object, server);
return "42";
});
Assert.DoesNotContain("server", JsonSerializer.Serialize(tool.ProtocolTool.InputSchema));
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Equal("42", result.Content[0].Text);
}
[Fact]
public async Task SupportsServiceFromDI()
{
MyService expectedMyService = new();
ServiceCollection sc = new();
sc.AddSingleton(expectedMyService);
IServiceProvider services = sc.BuildServiceProvider();
McpServerTool tool = McpServerTool.Create((MyService actualMyService) =>
{
Assert.Same(expectedMyService, actualMyService);
return "42";
}, services: services);
Assert.DoesNotContain("actualMyService", JsonSerializer.Serialize(tool.ProtocolTool.InputSchema));
Mock<IMcpServer> mockServer = new();
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.True(result.IsError);
mockServer.SetupGet(x => x.Services).Returns(services);
result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(mockServer.Object, null),
TestContext.Current.CancellationToken);
Assert.Equal("42", result.Content[0].Text);
}
[Fact]
public async Task SupportsOptionalServiceFromDI()
{
MyService expectedMyService = new();
ServiceCollection sc = new();
sc.AddSingleton(expectedMyService);
IServiceProvider services = sc.BuildServiceProvider();
McpServerTool tool = McpServerTool.Create((MyService? actualMyService = null) =>
{
Assert.Null(actualMyService);
return "42";
}, services: services);
var result = await tool.InvokeAsync(
new RequestContext<CallToolRequestParams>(null!, null),
TestContext.Current.CancellationToken);
Assert.Equal("42", result.Content[0].Text);
}
private sealed class MyService;
}

View file

@ -89,7 +89,7 @@ public class SseIntegrationTests
defaultOptions,
loggerFactory: loggerFactory,
cancellationToken: TestContext.Current.CancellationToken);
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken).ToListAsync(TestContext.Current.CancellationToken);
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
// assert
Assert.NotEmpty(tools);
@ -157,7 +157,7 @@ public class SseIntegrationTests
cancellationToken: TestContext.Current.CancellationToken);
// Call the server's sampleLLM tool which should trigger our sampling handler
var result = await client.CallToolAsync("sampleLLM", new Dictionary<string, object>
var result = await client.CallToolAsync("sampleLLM", new Dictionary<string, object?>
{
["prompt"] = "Test prompt",
["maxTokens"] = 100

View file

@ -56,7 +56,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
// act
var client = await GetClientAsync();
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken).ToListAsync(TestContext.Current.CancellationToken);
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
// assert
Assert.NotNull(tools);
@ -71,7 +71,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
var client = await GetClientAsync();
var result = await client.CallToolAsync(
"echo",
new Dictionary<string, object>
new Dictionary<string, object?>
{
["message"] = "Hello MCP!"
},
@ -93,15 +93,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
// act
var client = await GetClientAsync();
List<Resource> allResources = [];
string? cursor = null;
do
{
var resources = await client.ListResourcesAsync(cursor, CancellationToken.None);
allResources.AddRange(resources.Resources);
cursor = resources.NextCursor;
}
while (cursor != null);
IList<Resource> allResources = await client.ListResourcesAsync(TestContext.Current.CancellationToken);
// The everything server provides 100 test resources
Assert.Equal(100, allResources.Count);
@ -148,7 +140,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
// act
var client = await GetClientAsync();
var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken).ToListAsync(TestContext.Current.CancellationToken);
var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
// assert
Assert.NotNull(prompts);
@ -179,7 +171,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
// act
var client = await GetClientAsync();
var arguments = new Dictionary<string, object>
var arguments = new Dictionary<string, object?>
{
{ "temperature", "0.7" },
{ "style", "formal" }
@ -236,7 +228,7 @@ public class SseServerIntegrationTests : LoggedTest, IClassFixture<SseServerInte
#pragma warning restore CS1998 // Async method lacks 'await' operators and will run synchronously
// Call the server's sampleLLM tool which should trigger our sampling handler
var result = await client.CallToolAsync("sampleLLM", new Dictionary<string, object>
var result = await client.CallToolAsync("sampleLLM", new Dictionary<string, object?>
{
["prompt"] = "Test prompt",
["maxTokens"] = 100

View file

@ -0,0 +1,73 @@
using Microsoft.Extensions.Logging.Abstractions;
using ModelContextProtocol.Protocol.Messages;
using ModelContextProtocol.Protocol.Transport;
using ModelContextProtocol.Utils.Json;
using System.Text.Json;
namespace ModelContextProtocol.Tests.Transport;
internal sealed class StreamClientTransport : TransportBase, IClientTransport
{
private readonly JsonSerializerOptions _jsonOptions = McpJsonUtilities.DefaultOptions;
private Task? _readTask;
private CancellationTokenSource _shutdownCts = new CancellationTokenSource();
private readonly TextReader _stdin;
private readonly TextWriter _stdout;
public StreamClientTransport(TextReader stdin, TextWriter stdout)
: base(NullLoggerFactory.Instance)
{
_stdin = stdin;
_stdout = stdout;
_readTask = Task.Run(() => ReadMessagesAsync(_shutdownCts.Token), CancellationToken.None);
SetConnected(true);
}
public Task ConnectAsync(CancellationToken cancellationToken = default) => Task.CompletedTask;
public override async Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{
string id = message is IJsonRpcMessageWithId messageWithId ?
messageWithId.Id.ToString() :
"(no id)";
await _stdout.WriteLineAsync(JsonSerializer.Serialize(message)).ConfigureAwait(false);
await _stdout.FlushAsync(cancellationToken).ConfigureAwait(false);
}
private async Task ReadMessagesAsync(CancellationToken cancellationToken)
{
while (await _stdin.ReadLineAsync(cancellationToken).ConfigureAwait(false) is string line)
{
if (!string.IsNullOrWhiteSpace(line))
{
try
{
if (JsonSerializer.Deserialize<IJsonRpcMessage>(line.Trim(), _jsonOptions) is { } message)
{
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
}
}
catch (JsonException)
{
}
}
}
}
public override async ValueTask DisposeAsync()
{
if (_shutdownCts is { } shutdownCts)
{
await shutdownCts.CancelAsync().ConfigureAwait(false);
shutdownCts.Dispose();
}
if (_readTask is Task readTask)
{
await readTask.WaitAsync(TimeSpan.FromSeconds(5)).ConfigureAwait(false);
}
SetConnected(false);
}
}

View file

@ -76,15 +76,6 @@ public class TestServerTransport : IServerTransport
}, cancellationToken);
}
private async Task Error(JsonRpcRequest request, CancellationToken cancellationToken)
{
await WriteMessageAsync(new JsonRpcError
{
Id = request.Id,
Error = new JsonRpcErrorDetail() { Code = -32601, Message = $"Method '{request.Method}' not supported" }
}, cancellationToken);
}
protected async Task WriteMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default)
{
await _messageChannel.Writer.WriteAsync(message, cancellationToken);