Skip to content

Commit 4a53f6f

Browse files
halter73halter73
authored
List tools from DI once (modelcontextprotocol#115)
* List tools from DI once - Follow best practice to not mutate options outside of IConfigureOptions - Long-term, we should move the repeated work from SetToolsHandler into an IPostConfigureOptions service or IMcpServerConnectionFactory * Add Can_Create_Multiple_Servers_From_Options_And_List_Registered_Tools * Flip stdin and stdout in StreamClientTransport * s/McpToolType/McpServerToolType * Clarify stdin and stdout refers to server process in StreamClientTransport * fixup
1 parent 3c72c34 commit 4a53f6f

File tree

7 files changed

+133
-49
lines changed

7 files changed

+133
-49
lines changed

‎README.MD‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ var response = await chatClient.GetResponseAsync(
8080

8181
Here is an example of how to create an MCP server and register all tools from the current application.
8282
It includes a simple echo tool as an example (this is included in the same file here for easy of copy and paste, but it needn't be in the same file...
83-
the employed overload of `WithTools` examines the current assembly for classes with the `McpToolType` attribute, and registers all methods with the
83+
the employed overload of `WithTools` examines the current assembly for classes with the `McpServerToolType` attribute, and registers all methods with the
8484
`McpTool` attribute as tools.)
8585

8686
```csharp

‎src/ModelContextProtocol/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ var response = await chatClient.GetResponseAsync(
8585

8686
Here is an example of how to create an MCP server and register all tools from the current application.
8787
It includes a simple echo tool as an example (this is included in the same file here for easy of copy and paste, but it needn't be in the same file...
88-
the employed overload of `WithTools` examines the current assembly for classes with the `McpToolType` attribute, and registers all methods with the
88+
the employed overload of `WithTools` examines the current assembly for classes with the `McpServerToolType` attribute, and registers all methods with the
8989
`McpTool` attribute as tools.)
9090

9191
```csharp
@@ -101,7 +101,7 @@ builder.Services
101101
.WithTools();
102102
await builder.Build().RunAsync();
103103

104-
[McpToolType]
104+
[McpServerToolType]
105105
public static class EchoTool
106106
{
107107
[McpTool, Description("Echoes the message back to the client.")]

‎src/ModelContextProtocol/Server/McpServer.cs‎

Lines changed: 43 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
1515
{
1616
private readonly IServerTransport? _serverTransport;
1717
private readonly string _serverDescription;
18+
private readonly EventHandler? _toolsChangedDelegate;
19+
1820
private volatile bool _isInitializing;
1921

2022
/// <summary>
@@ -32,36 +34,45 @@ public McpServer(ITransport transport, McpServerOptions options, ILoggerFactory?
3234
Throw.IfNull(options);
3335

3436
_serverTransport = transport as IServerTransport;
35-
ServerInstructions = options.ServerInstructions;
37+
ServerOptions = options;
3638
Services = serviceProvider;
3739
_serverDescription = $"{options.ServerInfo.Name} {options.ServerInfo.Version}";
40+
_toolsChangedDelegate = delegate
41+
{
42+
_ = SendMessageAsync(new JsonRpcNotification()
43+
{
44+
Method = NotificationMethods.ToolListChangedNotification,
45+
});
46+
};
3847

3948
AddNotificationHandler("notifications/initialized", _ =>
4049
{
50+
if (ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
51+
{
52+
tools.Changed += _toolsChangedDelegate;
53+
}
54+
4155
IsInitialized = true;
4256
return Task.CompletedTask;
4357
});
4458

45-
SetToolsHandler(ref options);
59+
SetToolsHandler(options);
4660

4761
SetInitializeHandler(options);
4862
SetCompletionHandler(options);
4963
SetPingHandler();
5064
SetPromptsHandler(options);
5165
SetResourcesHandler(options);
5266
SetSetLoggingLevelHandler(options);
53-
54-
ServerOptions = options;
5567
}
5668

69+
public ServerCapabilities? ServerCapabilities { get; set; }
70+
5771
public ClientCapabilities? ClientCapabilities { get; set; }
5872

5973
/// <inheritdoc />
6074
public Implementation? ClientInfo { get; set; }
6175

62-
/// <inheritdoc />
63-
public string? ServerInstructions { get; set; }
64-
6576
/// <inheritdoc />
6677
public McpServerOptions ServerOptions { get; }
6778

@@ -111,6 +122,15 @@ public async Task StartAsync(CancellationToken cancellationToken = default)
111122
}
112123
}
113124

125+
protected override Task CleanupAsync()
126+
{
127+
if (ServerOptions.Capabilities?.Tools?.ToolCollection is { } tools)
128+
{
129+
tools.Changed -= _toolsChangedDelegate;
130+
}
131+
return base.CleanupAsync();
132+
}
133+
114134
private void SetPingHandler()
115135
{
116136
SetRequestHandler<JsonNode, PingResult>("ping",
@@ -127,9 +147,9 @@ private void SetInitializeHandler(McpServerOptions options)
127147
return Task.FromResult(new InitializeResult()
128148
{
129149
ProtocolVersion = options.ProtocolVersion,
130-
Instructions = ServerInstructions,
150+
Instructions = options.ServerInstructions,
131151
ServerInfo = options.ServerInfo,
132-
Capabilities = options.Capabilities ?? new ServerCapabilities(),
152+
Capabilities = ServerCapabilities ?? new(),
133153
});
134154
});
135155
}
@@ -198,7 +218,7 @@ private void SetPromptsHandler(McpServerOptions options)
198218
SetRequestHandler<GetPromptRequestParams, GetPromptResult>("prompts/get", (request, ct) => getPromptHandler(new(this, request), ct));
199219
}
200220

201-
private void SetToolsHandler(ref McpServerOptions options)
221+
private void SetToolsHandler(McpServerOptions options)
202222
{
203223
ToolsCapability? toolsCapability = options.Capabilities?.Tools;
204224
var listToolsHandler = toolsCapability?.ListToolsHandler;
@@ -261,25 +281,25 @@ private void SetToolsHandler(ref McpServerOptions options)
261281
return tool.InvokeAsync(request, cancellationToken);
262282
};
263283

264-
toolsCapability ??= new();
265-
toolsCapability.CallToolHandler = callToolHandler;
266-
toolsCapability.ListToolsHandler = listToolsHandler;
267-
toolsCapability.ToolCollection = tools;
268-
toolsCapability.ListChanged = true;
269-
270-
options.Capabilities ??= new();
271-
options.Capabilities.Tools = toolsCapability;
272-
273-
tools.Changed += delegate
284+
ServerCapabilities = new()
274285
{
275-
_ = SendMessageAsync(new JsonRpcNotification()
286+
Experimental = options.Capabilities?.Experimental,
287+
Logging = options.Capabilities?.Logging,
288+
Prompts = options.Capabilities?.Prompts,
289+
Resources = options.Capabilities?.Resources,
290+
Tools = new()
276291
{
277-
Method = NotificationMethods.ToolListChangedNotification,
278-
});
292+
ListToolsHandler = listToolsHandler,
293+
CallToolHandler = callToolHandler,
294+
ToolCollection = tools,
295+
ListChanged = true,
296+
}
279297
};
280298
}
281299
else
282300
{
301+
ServerCapabilities = options.Capabilities;
302+
283303
if (toolsCapability is null)
284304
{
285305
// No tools, and no tools capability was declared, so nothing to do.

‎src/ModelContextProtocol/Shared/McpJsonRpcEndpoint.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -359,7 +359,7 @@ protected void SetRequestHandler<TRequest, TResponse>(string method, Func<TReque
359359
/// Cleans up the endpoint and releases resources.
360360
/// </summary>
361361
/// <returns></returns>
362-
protected async Task CleanupAsync()
362+
protected virtual async Task CleanupAsync()
363363
{
364364
if (_isDisposed)
365365
return;

‎tests/ModelContextProtocol.Tests/Client/McpClientExtensionsTests.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,8 @@ private async Task<IMcpClient> CreateMcpClientForServer()
3838
{
3939
await _server.StartAsync(TestContext.Current.CancellationToken);
4040

41-
var stdin = new StreamReader(_serverToClientPipe.Reader.AsStream());
42-
var stdout = new StreamWriter(_clientToServerPipe.Writer.AsStream());
41+
var serverStdinWriter = new StreamWriter(_clientToServerPipe.Writer.AsStream());
42+
var serverStdoutReader = new StreamReader(_serverToClientPipe.Reader.AsStream());
4343

4444
var serverConfig = new McpServerConfig()
4545
{
@@ -50,7 +50,7 @@ private async Task<IMcpClient> CreateMcpClientForServer()
5050

5151
return await McpClientFactory.CreateAsync(
5252
serverConfig,
53-
createTransportFunc: (_, _) => new StreamClientTransport(stdin, stdout),
53+
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader),
5454
cancellationToken: TestContext.Current.CancellationToken);
5555
}
5656

‎tests/ModelContextProtocol.Tests/Configuration/McpServerBuilderExtensionsToolsTests.cs‎

Lines changed: 73 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -11,38 +11,45 @@
1111
using Microsoft.Extensions.AI;
1212
using System.Threading.Channels;
1313
using ModelContextProtocol.Protocol.Messages;
14+
using Microsoft.Extensions.Options;
15+
using ModelContextProtocol.Tests.Utils;
16+
using Microsoft.Extensions.Logging;
1417

1518
namespace ModelContextProtocol.Tests.Configuration;
1619

17-
public class McpServerBuilderExtensionsToolsTests : IAsyncDisposable
20+
public class McpServerBuilderExtensionsToolsTests : LoggedTest, IAsyncDisposable
1821
{
19-
private Pipe _clientToServerPipe = new();
20-
private Pipe _serverToClientPipe = new();
22+
private readonly Pipe _clientToServerPipe = new();
23+
private readonly Pipe _serverToClientPipe = new();
24+
private readonly ServiceProvider _serviceProvider;
2125
private readonly IMcpServerBuilder _builder;
2226
private readonly IMcpServer _server;
2327

24-
public McpServerBuilderExtensionsToolsTests()
28+
public McpServerBuilderExtensionsToolsTests(ITestOutputHelper testOutputHelper)
29+
: base(testOutputHelper)
2530
{
2631
ServiceCollection sc = new();
32+
sc.AddSingleton(LoggerFactory);
2733
sc.AddSingleton<IServerTransport>(new StdioServerTransport("TestServer", _clientToServerPipe.Reader.AsStream(), _serverToClientPipe.Writer.AsStream()));
2834
sc.AddSingleton(new ObjectWithId());
2935
_builder = sc.AddMcpServer().WithTools<EchoTool>();
30-
_server = sc.BuildServiceProvider().GetRequiredService<IMcpServer>();
36+
_serviceProvider = sc.BuildServiceProvider();
37+
_server = _serviceProvider.GetRequiredService<IMcpServer>();
3138
}
3239

3340
public ValueTask DisposeAsync()
3441
{
3542
_clientToServerPipe.Writer.Complete();
3643
_serverToClientPipe.Writer.Complete();
37-
return _server.DisposeAsync();
44+
return _serviceProvider.DisposeAsync();
3845
}
3946

4047
private async Task<IMcpClient> CreateMcpClientForServer()
4148
{
4249
await _server.StartAsync(TestContext.Current.CancellationToken);
4350

44-
var stdin = new StreamReader(_serverToClientPipe.Reader.AsStream());
45-
var stdout = new StreamWriter(_clientToServerPipe.Writer.AsStream());
51+
var serverStdinWriter = new StreamWriter(_clientToServerPipe.Writer.AsStream());
52+
var serverStdoutReader = new StreamReader(_serverToClientPipe.Reader.AsStream());
4653

4754
var serverConfig = new McpServerConfig()
4855
{
@@ -53,7 +60,7 @@ private async Task<IMcpClient> CreateMcpClientForServer()
5360

5461
return await McpClientFactory.CreateAsync(
5562
serverConfig,
56-
createTransportFunc: (_, _) => new StreamClientTransport(stdin, stdout),
63+
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader),
5764
cancellationToken: TestContext.Current.CancellationToken);
5865
}
5966

@@ -86,6 +93,63 @@ public async Task Can_List_Registered_Tools()
8693
Assert.Equal("Echoes the input back to the client.", doubleEchoTool.Description);
8794
}
8895

96+
97+
[Fact]
98+
public async Task Can_Create_Multiple_Servers_From_Options_And_List_Registered_Tools()
99+
{
100+
var options = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
101+
var loggerFactory = _serviceProvider.GetRequiredService<ILoggerFactory>();
102+
103+
for (int i = 0; i < 2; i++)
104+
{
105+
var stdinPipe = new Pipe();
106+
var stdoutPipe = new Pipe();
107+
108+
try
109+
{
110+
var transport = new StdioServerTransport($"TestServer_{i}", stdinPipe.Reader.AsStream(), stdoutPipe.Writer.AsStream());
111+
var server = McpServerFactory.Create(transport, options, loggerFactory, _serviceProvider);
112+
113+
await server.StartAsync(TestContext.Current.CancellationToken);
114+
115+
var serverStdinWriter = new StreamWriter(stdinPipe.Writer.AsStream());
116+
var serverStdoutReader = new StreamReader(stdoutPipe.Reader.AsStream());
117+
118+
var serverConfig = new McpServerConfig()
119+
{
120+
Id = $"TestServer_{i}",
121+
Name = $"TestServer_{i}",
122+
TransportType = "ignored",
123+
};
124+
125+
var client = await McpClientFactory.CreateAsync(
126+
serverConfig,
127+
createTransportFunc: (_, _) => new StreamClientTransport(serverStdinWriter, serverStdoutReader),
128+
cancellationToken: TestContext.Current.CancellationToken);
129+
130+
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
131+
Assert.Equal(11, tools.Count);
132+
133+
McpClientTool echoTool = tools.First(t => t.Name == "Echo");
134+
Assert.Equal("Echo", echoTool.Name);
135+
Assert.Equal("Echoes the input back to the client.", echoTool.Description);
136+
Assert.Equal("object", echoTool.JsonSchema.GetProperty("type").GetString());
137+
Assert.Equal(JsonValueKind.Object, echoTool.JsonSchema.GetProperty("properties").GetProperty("message").ValueKind);
138+
Assert.Equal("the echoes message", echoTool.JsonSchema.GetProperty("properties").GetProperty("message").GetProperty("description").GetString());
139+
Assert.Equal(1, echoTool.JsonSchema.GetProperty("required").GetArrayLength());
140+
141+
McpClientTool doubleEchoTool = tools.First(t => t.Name == "double_echo");
142+
Assert.Equal("double_echo", doubleEchoTool.Name);
143+
Assert.Equal("Echoes the input back to the client.", doubleEchoTool.Description);
144+
}
145+
finally
146+
{
147+
stdinPipe.Writer.Complete();
148+
stdoutPipe.Writer.Complete();
149+
}
150+
}
151+
}
152+
89153
[Fact]
90154
public async Task Can_Be_Notified_Of_Tool_Changes()
91155
{

‎tests/ModelContextProtocol.Tests/Transport/StreamClientTransport.cs‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -9,16 +9,16 @@ namespace ModelContextProtocol.Tests.Transport;
99
internal sealed class StreamClientTransport : TransportBase, IClientTransport
1010
{
1111
private readonly JsonSerializerOptions _jsonOptions = McpJsonUtilities.DefaultOptions;
12-
private Task? _readTask;
13-
private CancellationTokenSource _shutdownCts = new CancellationTokenSource();
14-
private readonly TextReader _stdin;
15-
private readonly TextWriter _stdout;
12+
private readonly Task? _readTask;
13+
private readonly CancellationTokenSource _shutdownCts = new CancellationTokenSource();
14+
private readonly TextReader _serverStdoutReader;
15+
private readonly TextWriter _serverStdinWriter;
1616

17-
public StreamClientTransport(TextReader stdin, TextWriter stdout)
17+
public StreamClientTransport(TextWriter serverStdinWriter, TextReader serverStdoutReader)
1818
: base(NullLoggerFactory.Instance)
1919
{
20-
_stdin = stdin;
21-
_stdout = stdout;
20+
_serverStdoutReader = serverStdoutReader;
21+
_serverStdinWriter = serverStdinWriter;
2222
_readTask = Task.Run(() => ReadMessagesAsync(_shutdownCts.Token), CancellationToken.None);
2323
SetConnected(true);
2424
}
@@ -31,13 +31,13 @@ public override async Task SendMessageAsync(IJsonRpcMessage message, Cancellatio
3131
messageWithId.Id.ToString() :
3232
"(no id)";
3333

34-
await _stdout.WriteLineAsync(JsonSerializer.Serialize(message)).ConfigureAwait(false);
35-
await _stdout.FlushAsync(cancellationToken).ConfigureAwait(false);
34+
await _serverStdinWriter.WriteLineAsync(JsonSerializer.Serialize(message)).ConfigureAwait(false);
35+
await _serverStdinWriter.FlushAsync(cancellationToken).ConfigureAwait(false);
3636
}
3737

3838
private async Task ReadMessagesAsync(CancellationToken cancellationToken)
3939
{
40-
while (await _stdin.ReadLineAsync(cancellationToken).ConfigureAwait(false) is string line)
40+
while (await _serverStdoutReader.ReadLineAsync(cancellationToken).ConfigureAwait(false) is string line)
4141
{
4242
if (!string.IsNullOrWhiteSpace(line))
4343
{

0 commit comments

Comments
 (0)