Skip to content

Commit ab5a950

Browse files
Tyler-R-KendrickTyler-R-Kendrick
committed
removed addnotificationhandler from client.
1 parent 6114cb9 commit ab5a950

File tree

9 files changed

+156
-81
lines changed

9 files changed

+156
-81
lines changed

‎src/ModelContextProtocol/Client/IMcpClient.cs‎

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
using ModelContextProtocol.Protocol.Types;
22
using ModelContextProtocol.Shared;
3-
using ModelContextProtocol.Protocol.Messages;
43

54
namespace ModelContextProtocol.Client;
65

@@ -25,21 +24,4 @@ public interface IMcpClient : IMcpSession
2524
/// It can be thought of like a "hint" to the model. For example, this information MAY be added to the system prompt.
2625
/// </summary>
2726
string? ServerInstructions { get; }
28-
29-
30-
/// <summary>
31-
/// Adds a handler for server notifications of a specific method.
32-
/// </summary>
33-
/// <param name="method">The notification method to handle.</param>
34-
/// <param name="handler">The async handler function to process notifications.</param>
35-
/// <remarks>
36-
/// <para>
37-
/// Each method may have multiple handlers. Adding a handler for a method that already has one
38-
/// will not replace the existing handler.
39-
/// </para>
40-
/// <para>
41-
/// <see cref="NotificationMethods"> provides constants for common notification methods.</see>
42-
/// </para>
43-
/// </remarks>
44-
void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
4527
}

‎src/ModelContextProtocol/Client/McpClient.cs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@ public McpClient(IClientTransport clientTransport, McpClientOptions options, Mcp
5959
RequestMethods.RootsList,
6060
(request, cancellationToken) => rootsHandler(request, cancellationToken));
6161
}
62+
63+
SetNotificationHandlers(options.NotificationHandlers);
6264
}
6365

6466
/// <inheritdoc/>
@@ -141,6 +143,20 @@ await SendMessageAsync(
141143
}
142144
}
143145

146+
private void SetNotificationHandlers(
147+
IReadOnlyDictionary<string, List<Func<JsonRpcNotification, Task>>> notificationHandlers)
148+
{
149+
foreach (var handlers in notificationHandlers)
150+
{
151+
var key = handlers.Key;
152+
var list = handlers.Value;
153+
foreach (var item in list)
154+
{
155+
AddNotificationHandler(key, item);
156+
}
157+
}
158+
}
159+
144160
/// <inheritdoc/>
145161
public override async ValueTask DisposeUnsynchronizedAsync()
146162
{

‎src/ModelContextProtocol/Client/McpClientOptions.cs‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,7 @@
1-
using ModelContextProtocol.Protocol.Types;
1+
using System.Text.Json.Serialization;
2+
using ModelContextProtocol.Protocol.Types;
3+
using ModelContextProtocol.Protocol.Messages;
4+
using ModelContextProtocol.Shared;
25

36
namespace ModelContextProtocol.Client;
47

@@ -28,4 +31,11 @@ public class McpClientOptions
2831
/// Timeout for initialization sequence.
2932
/// </summary>
3033
public TimeSpan InitializationTimeout { get; set; } = TimeSpan.FromSeconds(60);
34+
35+
/// <summary>
36+
/// Gets or sets the handler for get notifications.
37+
/// </summary>
38+
[JsonIgnore]
39+
public IReadOnlyDictionary<string, List<Func<JsonRpcNotification, Task>>> NotificationHandlers { get; init; } = new NotificationHandlers();
40+
3141
}

‎src/ModelContextProtocol/Server/McpServer.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ internal sealed class McpServer : McpJsonRpcEndpoint, IMcpServer
1414
private readonly EventHandler? _toolsChangedDelegate;
1515
private readonly EventHandler? _promptsChangedDelegate;
1616

17-
private ITransport _sessionTransport;
17+
private readonly ITransport _sessionTransport;
1818
private string _endpointName;
1919

2020
/// <summary>

‎src/ModelContextProtocol/Server/McpServerOptions.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ public class McpServerOptions
4444
public Func<RequestContext<CompleteRequestParams>, CancellationToken, Task<CompleteResult>>? GetCompletionHandler { get; set; }
4545

4646
/// <summary>
47-
/// Gets or sets the handler for get completion requests.
47+
/// Gets or sets the handler for get notifications.
4848
/// </summary>
4949
[JsonIgnore]
5050
public IReadOnlyDictionary<string, List<Func<JsonRpcNotification, Task>>> NotificationHandlers { get; init; } = new NotificationHandlers();

‎tests/ModelContextProtocol.Tests/ClientIntegrationTests.cs‎

Lines changed: 49 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -253,15 +253,26 @@ public async Task SubscribeResource_Stdio()
253253
{
254254
// arrange
255255
var clientId = "test_server";
256-
257-
// act
258256
TaskCompletionSource<bool> tcs = new();
259-
await using var client = await _fixture.CreateClientAsync(clientId);
260-
client.AddNotificationHandler(NotificationMethods.ResourceUpdatedNotification, (notification) =>
257+
Task HandleResourceUpdatedNotification(JsonRpcNotification notification)
261258
{
262259
var notificationParams = JsonSerializer.Deserialize<ResourceUpdatedNotificationParams>(notification.Params!.ToString() ?? string.Empty);
263260
tcs.TrySetResult(true);
264261
return Task.CompletedTask;
262+
}
263+
264+
// act
265+
await using var client = await _fixture.CreateClientAsync(clientId, new()
266+
{
267+
ClientInfo = new()
268+
{
269+
Name = "IntegrationTestClient",
270+
Version = "1.0.0"
271+
},
272+
NotificationHandlers = new Dictionary<string, List<Func<JsonRpcNotification, Task>>>()
273+
{
274+
[NotificationMethods.ResourceUpdatedNotification] = [HandleResourceUpdatedNotification],
275+
},
265276
});
266277
await client.SubscribeToResourceAsync("test://static/resource/1", TestContext.Current.CancellationToken);
267278

@@ -274,15 +285,26 @@ public async Task UnsubscribeResource_Stdio()
274285
{
275286
// arrange
276287
var clientId = "test_server";
277-
278-
// act
279288
TaskCompletionSource<bool> receivedNotification = new();
280-
await using var client = await _fixture.CreateClientAsync(clientId);
281-
client.AddNotificationHandler(NotificationMethods.ResourceUpdatedNotification, (notification) =>
289+
Task HandleResourceUpdatedNotification(JsonRpcNotification notification)
282290
{
283291
var notificationParams = JsonSerializer.Deserialize<ResourceUpdatedNotificationParams>(notification.Params!.ToString() ?? string.Empty);
284292
receivedNotification.TrySetResult(true);
285293
return Task.CompletedTask;
294+
}
295+
296+
// act
297+
await using var client = await _fixture.CreateClientAsync(clientId, new()
298+
{
299+
ClientInfo = new()
300+
{
301+
Name = "IntegrationTestClient",
302+
Version = "1.0.0"
303+
},
304+
NotificationHandlers = new Dictionary<string, List<Func<JsonRpcNotification, Task>>>()
305+
{
306+
[NotificationMethods.ResourceUpdatedNotification] = [HandleResourceUpdatedNotification],
307+
},
286308
});
287309
await client.SubscribeToResourceAsync("test://static/resource/1", TestContext.Current.CancellationToken);
288310

@@ -543,6 +565,16 @@ public async Task SamplingViaChatClient_RequestResponseProperlyPropagated()
543565
public async Task SetLoggingLevel_ReceivesLoggingMessages(string clientId)
544566
{
545567
// arrange
568+
TaskCompletionSource<bool> receivedNotification = new();
569+
Task HandleLoggingNotification(JsonRpcNotification notification)
570+
{
571+
var loggingMessageNotificationParameters = JsonSerializer.Deserialize<LoggingMessageNotificationParams>(notification.Params!.ToString() ?? string.Empty);
572+
if (loggingMessageNotificationParameters is not null)
573+
{
574+
receivedNotification.TrySetResult(true);
575+
}
576+
return Task.CompletedTask;
577+
}
546578
JsonSerializerOptions jsonSerializerOptions = new(JsonSerializerDefaults.Web)
547579
{
548580
TypeInfoResolver = new DefaultJsonTypeInfoResolver(),
@@ -552,17 +584,17 @@ public async Task SetLoggingLevel_ReceivesLoggingMessages(string clientId)
552584
Encoder = JavaScriptEncoder.UnsafeRelaxedJsonEscaping,
553585
};
554586

555-
TaskCompletionSource<bool> receivedNotification = new();
556-
await using var client = await _fixture.CreateClientAsync(clientId);
557-
client.AddNotificationHandler(NotificationMethods.LoggingMessageNotification, (notification) =>
587+
await using var client = await _fixture.CreateClientAsync(clientId, new()
558588
{
559-
var loggingMessageNotificationParameters = JsonSerializer.Deserialize<LoggingMessageNotificationParams>(notification.Params!.ToString() ?? string.Empty,
560-
jsonSerializerOptions);
561-
if (loggingMessageNotificationParameters is not null)
589+
ClientInfo = new()
562590
{
563-
receivedNotification.TrySetResult(true);
564-
}
565-
return Task.CompletedTask;
591+
Name = "IntegrationTestClient",
592+
Version = "1.0.0"
593+
},
594+
NotificationHandlers = new Dictionary<string, List<Func<JsonRpcNotification, Task>>>()
595+
{
596+
[NotificationMethods.LoggingMessageNotification] = [HandleLoggingNotification],
597+
},
566598
});
567599

568600
// act

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

Lines changed: 27 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,8 @@ public async ValueTask DisposeAsync()
117117
Dispose();
118118
}
119119

120-
private async Task<IMcpClient> CreateMcpClientForServer()
120+
private async Task<IMcpClient> CreateMcpClientForServer(
121+
Dictionary<string, List<Func<JsonRpcNotification, Task>>>? notificationHandlers = null)
121122
{
122123
return await McpClientFactory.CreateAsync(
123124
new McpServerConfig()
@@ -126,6 +127,15 @@ private async Task<IMcpClient> CreateMcpClientForServer()
126127
Name = "TestServer",
127128
TransportType = "ignored",
128129
},
130+
clientOptions: new()
131+
{
132+
ClientInfo = new()
133+
{
134+
Name = "TestClient",
135+
Version = "1.0.0",
136+
},
137+
NotificationHandlers = notificationHandlers ?? [],
138+
},
129139
createTransportFunc: (_, _) => new StreamClientTransport(
130140
serverInput: _clientToServerPipe.Writer.AsStream(),
131141
serverOutput: _serverToClientPipe.Reader.AsStream(),
@@ -175,19 +185,24 @@ public async Task Can_List_And_Call_Registered_Prompts()
175185
[Fact]
176186
public async Task Can_Be_Notified_Of_Prompt_Changes()
177187
{
178-
IMcpClient client = await CreateMcpClientForServer();
179-
180-
var prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
181-
Assert.Equal(6, prompts.Count);
182-
188+
var token = TestContext.Current.CancellationToken;
183189
Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
184-
client.AddNotificationHandler("notifications/prompts/list_changed", notification =>
190+
Task HandleListChangedNotification(JsonRpcNotification notification)
185191
{
186192
listChanged.Writer.TryWrite(notification);
187193
return Task.CompletedTask;
188-
});
194+
}
195+
IMcpClient client = await CreateMcpClientForServer(
196+
notificationHandlers: new()
197+
{
198+
[NotificationMethods.PromptListChangedNotification] = [HandleListChangedNotification],
199+
}
200+
);
201+
202+
var prompts = await client.ListPromptsAsync(token);
203+
Assert.Equal(6, prompts.Count);
189204

190-
var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
205+
var notificationRead = listChanged.Reader.ReadAsync(token);
191206
Assert.False(notificationRead.IsCompleted);
192207

193208
var serverOptions = _serviceProvider.GetRequiredService<IOptions<McpServerOptions>>().Value;
@@ -198,16 +213,16 @@ public async Task Can_Be_Notified_Of_Prompt_Changes()
198213
serverPrompts.Add(newPrompt);
199214
await notificationRead;
200215

201-
prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
216+
prompts = await client.ListPromptsAsync(token);
202217
Assert.Equal(7, prompts.Count);
203218
Assert.Contains(prompts, t => t.Name == "NewPrompt");
204219

205-
notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
220+
notificationRead = listChanged.Reader.ReadAsync(token);
206221
Assert.False(notificationRead.IsCompleted);
207222
serverPrompts.Remove(newPrompt);
208223
await notificationRead;
209224

210-
prompts = await client.ListPromptsAsync(TestContext.Current.CancellationToken);
225+
prompts = await client.ListPromptsAsync(token);
211226
Assert.Equal(6, prompts.Count);
212227
Assert.DoesNotContain(prompts, t => t.Name == "NewPrompt");
213228
}

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

Lines changed: 38 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -141,15 +141,25 @@ public async ValueTask DisposeAsync()
141141
Dispose();
142142
}
143143

144-
private async Task<IMcpClient> CreateMcpClientForServer()
144+
private async Task<IMcpClient> CreateMcpClientForServer(
145+
IReadOnlyDictionary<string, List<Func<JsonRpcNotification, Task>>>? notificationHandlers = null)
145146
{
146147
return await McpClientFactory.CreateAsync(
147-
new McpServerConfig()
148-
{
149-
Id = "TestServer",
150-
Name = "TestServer",
151-
TransportType = "ignored",
152-
},
148+
new McpServerConfig()
149+
{
150+
Id = "TestServer",
151+
Name = "TestServer",
152+
TransportType = "ignored",
153+
},
154+
clientOptions: new()
155+
{
156+
ClientInfo = new()
157+
{
158+
Name = "TestClient",
159+
Version = "1.0.0",
160+
},
161+
NotificationHandlers = notificationHandlers ?? new Dictionary<string, List<Func<JsonRpcNotification, Task>>>(),
162+
},
153163
createTransportFunc: (_, _) => new StreamClientTransport(
154164
serverInput: _clientToServerPipe.Writer.AsStream(),
155165
_serverToClientPipe.Reader.AsStream(),
@@ -242,17 +252,21 @@ public async Task Can_Create_Multiple_Servers_From_Options_And_List_Registered_T
242252
[Fact]
243253
public async Task Can_Be_Notified_Of_Tool_Changes()
244254
{
245-
IMcpClient client = await CreateMcpClientForServer();
246-
247-
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
248-
Assert.Equal(16, tools.Count);
249-
250255
Channel<JsonRpcNotification> listChanged = Channel.CreateUnbounded<JsonRpcNotification>();
251-
client.AddNotificationHandler(NotificationMethods.ToolListChangedNotification, notification =>
256+
Task HandleListChangeNotification(JsonRpcNotification notification)
252257
{
253258
listChanged.Writer.TryWrite(notification);
254259
return Task.CompletedTask;
255-
});
260+
}
261+
IMcpClient client = await CreateMcpClientForServer(
262+
notificationHandlers: new Dictionary<string, List<Func<JsonRpcNotification, Task>>>()
263+
{
264+
[NotificationMethods.ToolListChangedNotification] = [HandleListChangeNotification],
265+
}
266+
);
267+
268+
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
269+
Assert.Equal(16, tools.Count);
256270

257271
var notificationRead = listChanged.Reader.ReadAsync(TestContext.Current.CancellationToken);
258272
Assert.False(notificationRead.IsCompleted);
@@ -607,17 +621,21 @@ public void Create_ExtractsToolAnnotations_SomeSet()
607621
[Fact]
608622
public async Task HandlesIProgressParameter()
609623
{
624+
var token = TestContext.Current.CancellationToken;
610625
ConcurrentQueue<ProgressNotification> notifications = new();
611-
612-
IMcpClient client = await CreateMcpClientForServer();
613-
client.AddNotificationHandler(NotificationMethods.ProgressNotification, notification =>
626+
Task HandleProgressNotification(JsonRpcNotification notification)
614627
{
615628
ProgressNotification pn = JsonSerializer.Deserialize<ProgressNotification>((JsonElement)notification.Params!)!;
616629
notifications.Enqueue(pn);
617630
return Task.CompletedTask;
618-
});
631+
}
632+
IMcpClient client = await CreateMcpClientForServer(
633+
notificationHandlers: new Dictionary<string, List<Func<JsonRpcNotification, Task>>>()
634+
{
635+
[NotificationMethods.ProgressNotification] = [HandleProgressNotification],
636+
});
619637

620-
var tools = await client.ListToolsAsync(TestContext.Current.CancellationToken);
638+
var tools = await client.ListToolsAsync(token);
621639
Assert.NotNull(tools);
622640
Assert.NotEmpty(tools);
623641

@@ -631,7 +649,7 @@ public async Task HandlesIProgressParameter()
631649
Name = progressTool.ProtocolTool.Name,
632650
Meta = new() { ProgressToken = new("abc123") },
633651
},
634-
}, TestContext.Current.CancellationToken);
652+
}, token);
635653

636654
Assert.Contains("done", JsonSerializer.Serialize(result));
637655
SpinWait.SpinUntil(() => notifications.Count == 10, TimeSpan.FromSeconds(10));

0 commit comments

Comments
 (0)