Skip to content

Commit 4584f2c

Browse files
stephentoubstephentoub
authored
Merge branch 'main' into stvansolano/add-more-tests
2 parents 52c9a47 + 674cb15 commit 4584f2c

30 files changed

+461
-299
lines changed

‎src/ModelContextProtocol.AspNetCore/McpEndpointRouteBuilderExtensions.cs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,11 +47,11 @@ public static IEndpointConventionBuilder MapMcp(this IEndpointRouteBuilder endpo
4747
{
4848
throw new Exception($"Unreachable given good entropy! Session with ID '{sessionId}' has already been created.");
4949
}
50-
await using var server = McpServerFactory.Create(transport, mcpServerOptions.Value, loggerFactory, endpoints.ServiceProvider);
5150

5251
try
5352
{
5453
var transportTask = transport.RunAsync(cancellationToken: requestAborted);
54+
await using var server = McpServerFactory.Create(transport, mcpServerOptions.Value, loggerFactory, endpoints.ServiceProvider);
5555

5656
try
5757
{
@@ -85,7 +85,7 @@ public static IEndpointConventionBuilder MapMcp(this IEndpointRouteBuilder endpo
8585

8686
if (!_sessions.TryGetValue(sessionId.ToString(), out var transport))
8787
{
88-
await Results.BadRequest($"Session {sessionId} not found.").ExecuteAsync(context);
88+
await Results.BadRequest($"Session ID not found.").ExecuteAsync(context);
8989
return;
9090
}
9191

‎src/ModelContextProtocol/Client/IMcpClient.cs‎

Lines changed: 2 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,11 @@
1-
using ModelContextProtocol.Protocol.Messages;
2-
using ModelContextProtocol.Protocol.Types;
1+
using ModelContextProtocol.Protocol.Types;
32

43
namespace ModelContextProtocol.Client;
54

65
/// <summary>
76
/// Represents an instance of an MCP client connecting to a specific server.
87
/// </summary>
9-
public interface IMcpClient : IAsyncDisposable
8+
public interface IMcpClient : IMcpEndpoint
109
{
1110
/// <summary>
1211
/// Gets the capabilities supported by the server.
@@ -24,40 +23,4 @@ public interface IMcpClient : IAsyncDisposable
2423
/// It can be thought of like a "hint" to the model. For example, this information MAY be added to the system prompt.
2524
/// </summary>
2625
string? ServerInstructions { get; }
27-
28-
/// <summary>
29-
/// Adds a handler for server notifications of a specific method.
30-
/// </summary>
31-
/// <param name="method">The notification method to handle.</param>
32-
/// <param name="handler">The async handler function to process notifications.</param>
33-
/// <remarks>
34-
/// <para>
35-
/// Each method may have multiple handlers. Adding a handler for a method that already has one
36-
/// will not replace the existing handler.
37-
/// </para>
38-
/// <para>
39-
/// <see cref="NotificationMethods"> provides constants for common notification methods.</see>
40-
/// </para>
41-
/// </remarks>
42-
void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
43-
44-
/// <summary>
45-
/// Sends a generic JSON-RPC request to the server.
46-
/// </summary>
47-
/// <typeparam name="TResult">The expected response type.</typeparam>
48-
/// <param name="request">The JSON-RPC request to send.</param>
49-
/// <param name="cancellationToken">A token to cancel the operation.</param>
50-
/// <returns>A task containing the server's response.</returns>
51-
/// <remarks>
52-
/// It is recommended to use the capability-specific methods that use this one in their implementation.
53-
/// Use this method for custom requests or those not yet covered explicitly.
54-
/// </remarks>
55-
Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class;
56-
57-
/// <summary>
58-
/// Sends a message to the server.
59-
/// </summary>
60-
/// <param name="message">The message.</param>
61-
/// <param name="cancellationToken">A token to cancel the operation.</param>
62-
Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
6326
}

‎src/ModelContextProtocol/Client/McpClient.cs‎

Lines changed: 20 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,10 @@ public McpClient(IClientTransport clientTransport, McpClientOptions options, Mcp
4242

4343
SetRequestHandler<CreateMessageRequestParams, CreateMessageResult>(
4444
RequestMethods.SamplingCreateMessage,
45-
(request, ct) => samplingHandler(request, ct));
45+
(request, cancellationToken) => samplingHandler(
46+
request,
47+
request?.Meta?.ProgressToken is { } token ? new TokenProgress(this, token) : NullProgress.Instance,
48+
cancellationToken));
4649
}
4750

4851
if (options.Capabilities?.Roots is { } rootsCapability)
@@ -54,7 +57,7 @@ public McpClient(IClientTransport clientTransport, McpClientOptions options, Mcp
5457

5558
SetRequestHandler<ListRootsRequestParams, ListRootsResult>(
5659
RequestMethods.RootsList,
57-
(request, ct) => rootsHandler(request, ct));
60+
(request, cancellationToken) => rootsHandler(request, cancellationToken));
5861
}
5962
}
6063

@@ -79,29 +82,27 @@ public async Task ConnectAsync(CancellationToken cancellationToken = default)
7982
{
8083
// Connect transport
8184
_sessionTransport = await _clientTransport.ConnectAsync(cancellationToken).ConfigureAwait(false);
82-
// We don't want the ConnectAsync token to cancel the session after we've successfully connected.
83-
// The base class handles cleaning up the session in DisposeAsync without our help.
84-
StartSession(_sessionTransport, fullSessionCancellationToken: CancellationToken.None);
85+
StartSession(_sessionTransport);
8586

8687
// Perform initialization sequence
8788
using var initializationCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
8889
initializationCts.CancelAfter(_options.InitializationTimeout);
8990

90-
try
91-
{
92-
// Send initialize request
93-
var initializeResponse = await SendRequestAsync<InitializeResult>(
94-
new JsonRpcRequest
95-
{
96-
Method = RequestMethods.Initialize,
97-
Params = new InitializeRequestParams()
91+
try
92+
{
93+
// Send initialize request
94+
var initializeResponse = await SendRequestAsync<InitializeResult>(
95+
new JsonRpcRequest
9896
{
99-
ProtocolVersion = _options.ProtocolVersion,
100-
Capabilities = _options.Capabilities ?? new ClientCapabilities(),
101-
ClientInfo = _options.ClientInfo
102-
}
103-
},
104-
initializationCts.Token).ConfigureAwait(false);
97+
Method = RequestMethods.Initialize,
98+
Params = new InitializeRequestParams()
99+
{
100+
ProtocolVersion = _options.ProtocolVersion,
101+
Capabilities = _options.Capabilities ?? new ClientCapabilities(),
102+
ClientInfo = _options.ClientInfo
103+
}
104+
},
105+
initializationCts.Token).ConfigureAwait(false);
105106

106107
// Store server information
107108
_logger.ServerCapabilitiesReceived(EndpointName,

‎src/ModelContextProtocol/Client/McpClientExtensions.cs‎

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,7 @@
88

99
namespace ModelContextProtocol.Client;
1010

11-
/// <summary>
12-
/// Provides extensions for operating on MCP clients.
13-
/// </summary>
11+
/// <summary>Provides extension methods for interacting with an <see cref="IMcpClient"/>.</summary>
1412
public static class McpClientExtensions
1513
{
1614
/// <summary>
@@ -531,17 +529,33 @@ internal static CreateMessageResult ToCreateMessageResult(this ChatResponse chat
531529
/// </summary>
532530
/// <param name="chatClient">The <see cref="IChatClient"/> with which to satisfy sampling requests.</param>
533531
/// <returns>The created handler delegate.</returns>
534-
public static Func<CreateMessageRequestParams?, CancellationToken, Task<CreateMessageResult>> CreateSamplingHandler(this IChatClient chatClient)
532+
public static Func<CreateMessageRequestParams?, IProgress<ProgressNotificationValue>, CancellationToken, Task<CreateMessageResult>> CreateSamplingHandler(
533+
this IChatClient chatClient)
535534
{
536535
Throw.IfNull(chatClient);
537536

538-
return async (requestParams, cancellationToken) =>
537+
return async (requestParams, progress, cancellationToken) =>
539538
{
540539
Throw.IfNull(requestParams);
541540

542541
var (messages, options) = requestParams.ToChatClientArguments();
543-
var response = await chatClient.GetResponseAsync(messages, options, cancellationToken).ConfigureAwait(false);
544-
return response.ToCreateMessageResult();
542+
var progressToken = requestParams.Meta?.ProgressToken;
543+
544+
List<ChatResponseUpdate> updates = [];
545+
await foreach (var update in chatClient.GetStreamingResponseAsync(messages, options, cancellationToken))
546+
{
547+
updates.Add(update);
548+
549+
if (progressToken is not null)
550+
{
551+
progress.Report(new()
552+
{
553+
Progress = updates.Count,
554+
});
555+
}
556+
}
557+
558+
return updates.ToChatResponse().ToCreateMessageResult();
545559
};
546560
}
547561

‎src/ModelContextProtocol/IMcpEndpoint.cs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
using ModelContextProtocol.Protocol.Messages;
2+
3+
namespace ModelContextProtocol;
4+
5+
/// <summary>Represents a client or server MCP endpoint.</summary>
6+
public interface IMcpEndpoint : IAsyncDisposable
7+
{
8+
/// <summary>Sends a generic JSON-RPC request to the connected endpoint.</summary>
9+
/// <typeparam name="TResult">The expected response type.</typeparam>
10+
/// <param name="request">The JSON-RPC request to send.</param>
11+
/// <param name="cancellationToken">A token to cancel the operation.</param>
12+
/// <returns>A task containing the client's response.</returns>
13+
Task<TResult> SendRequestAsync<TResult>(JsonRpcRequest request, CancellationToken cancellationToken = default) where TResult : class;
14+
15+
/// <summary>Sends a message to the connected endpoint.</summary>
16+
/// <param name="message">The message.</param>
17+
/// <param name="cancellationToken">A token to cancel the operation.</param>
18+
Task SendMessageAsync(IJsonRpcMessage message, CancellationToken cancellationToken = default);
19+
20+
/// <summary>
21+
/// Adds a handler for server notifications of a specific method.
22+
/// </summary>
23+
/// <param name="method">The notification method to handle.</param>
24+
/// <param name="handler">The async handler function to process notifications.</param>
25+
/// <remarks>
26+
/// <para>
27+
/// Each method may have multiple handlers. Adding a handler for a method that already has one
28+
/// will not replace the existing handler.
29+
/// </para>
30+
/// <para>
31+
/// <see cref="NotificationMethods"> provides constants for common notification methods.</see>
32+
/// </para>
33+
/// </remarks>
34+
void AddNotificationHandler(string method, Func<JsonRpcNotification, Task> handler);
35+
}

‎src/ModelContextProtocol/McpEndpointExtensions.cs‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
using ModelContextProtocol.Protocol.Messages;
2+
using ModelContextProtocol.Utils;
3+
4+
namespace ModelContextProtocol;
5+
6+
/// <summary>Provides extension methods for interacting with an <see cref="IMcpEndpoint"/>.</summary>
7+
public static class McpEndpointExtensions
8+
{
9+
/// <summary>Notifies the connected endpoint of progress.</summary>
10+
/// <param name="endpoint">The endpoint issueing the notification.</param>
11+
/// <param name="progressToken">The <see cref="ProgressToken"/> identifying the operation.</param>
12+
/// <param name="progress">The progress update to send.</param>
13+
/// <param name="cancellationToken">A token to cancel the operation.</param>
14+
/// <returns>A task representing the completion of the operation.</returns>
15+
/// <exception cref="ArgumentNullException"><paramref name="endpoint"/> is <see langword="null"/>.</exception>
16+
public static Task NotifyProgressAsync(
17+
this IMcpEndpoint endpoint,
18+
ProgressToken progressToken,
19+
ProgressNotificationValue progress,
20+
CancellationToken cancellationToken = default)
21+
{
22+
Throw.IfNull(endpoint);
23+
24+
return endpoint.SendMessageAsync(new JsonRpcNotification()
25+
{
26+
Method = NotificationMethods.ProgressNotification,
27+
Params = new ProgressNotification()
28+
{
29+
ProgressToken = progressToken,
30+
Progress = progress,
31+
},
32+
}, cancellationToken);
33+
}
34+
}

‎src/ModelContextProtocol/Protocol/Transport/SseResponseStreamTransport.cs‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -32,30 +32,30 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string
3232
/// <returns>A task representing the send loop that writes JSON-RPC messages to the SSE response stream.</returns>
3333
public Task RunAsync(CancellationToken cancellationToken)
3434
{
35-
void WriteJsonRpcMessageToBuffer(SseItem<IJsonRpcMessage?> item, IBufferWriter<byte> writer)
36-
{
37-
if (item.EventType == "endpoint")
38-
{
39-
writer.Write(Encoding.UTF8.GetBytes(messageEndpoint));
40-
return;
41-
}
42-
43-
JsonSerializer.Serialize(GetUtf8JsonWriter(writer), item.Data, McpJsonUtilities.DefaultOptions.GetTypeInfo<IJsonRpcMessage?>());
44-
}
45-
46-
IsConnected = true;
47-
4835
// The very first SSE event isn't really an IJsonRpcMessage, but there's no API to write a single item of a different type,
4936
// so we fib and special-case the "endpoint" event type in the formatter.
5037
if (!_outgoingSseChannel.Writer.TryWrite(new SseItem<IJsonRpcMessage?>(null, "endpoint")))
5138
{
5239
throw new InvalidOperationException($"You must call ${nameof(RunAsync)} before calling ${nameof(SendMessageAsync)}.");
5340
}
5441

42+
IsConnected = true;
43+
5544
var sseItems = _outgoingSseChannel.Reader.ReadAllAsync(cancellationToken);
5645
return _sseWriteTask = SseFormatter.WriteAsync(sseItems, sseResponseStream, WriteJsonRpcMessageToBuffer, cancellationToken);
5746
}
5847

48+
private void WriteJsonRpcMessageToBuffer(SseItem<IJsonRpcMessage?> item, IBufferWriter<byte> writer)
49+
{
50+
if (item.EventType == "endpoint")
51+
{
52+
writer.Write(Encoding.UTF8.GetBytes(messageEndpoint));
53+
return;
54+
}
55+
56+
JsonSerializer.Serialize(GetUtf8JsonWriter(writer), item.Data, McpJsonUtilities.DefaultOptions.GetTypeInfo<IJsonRpcMessage?>());
57+
}
58+
5959
/// <inheritdoc/>
6060
public ChannelReader<IJsonRpcMessage> MessageReader => _incomingChannel.Reader;
6161

‎src/ModelContextProtocol/Protocol/Types/Capabilities.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ public class SamplingCapability
5555

5656
/// <summary>Gets or sets the handler for sampling requests.</summary>
5757
[JsonIgnore]
58-
public Func<CreateMessageRequestParams?, CancellationToken, Task<CreateMessageResult>>? SamplingHandler { get; set; }
58+
public Func<CreateMessageRequestParams?, IProgress<ProgressNotificationValue>, CancellationToken, Task<CreateMessageResult>>? SamplingHandler { get; set; }
5959
}
6060

6161
/// <summary>

‎src/ModelContextProtocol/Protocol/Types/ListPromptsRequestParams.cs‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,4 @@
44
/// Sent from the client to request a list of prompts and prompt templates the server has.
55
/// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/">See the schema for details</see>
66
/// </summary>
7-
public class ListPromptsRequestParams
8-
{
9-
/// <summary>
10-
/// An opaque token representing the current pagination position.
11-
/// If provided, the server should return results starting after this cursor.
12-
/// </summary>
13-
[System.Text.Json.Serialization.JsonPropertyName("cursor")]
14-
public string? Cursor { get; init; }
15-
}
7+
public class ListPromptsRequestParams : PaginatedRequestParams;

‎src/ModelContextProtocol/Protocol/Types/ListResourceTemplatesRequestParams.cs‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,4 @@
44
/// Sent from the client to request a list of resource templates the server has.
55
/// <see href="https://github.com/modelcontextprotocol/specification/blob/main/schema/">See the schema for details</see>
66
/// </summary>
7-
public class ListResourceTemplatesRequestParams
8-
{
9-
/// <summary>
10-
/// An opaque token representing the current pagination position.
11-
/// If provided, the server should return results starting after this cursor.
12-
/// </summary>
13-
[System.Text.Json.Serialization.JsonPropertyName("cursor")]
14-
public string? Cursor { get; init; }
15-
}
7+
public class ListResourceTemplatesRequestParams : PaginatedRequestParams;

0 commit comments

Comments
 (0)