Pass session id's to MCP endpoints. (#466)
* Pass session id's to MCP endpoints. * Update src/ModelContextProtocol/Client/McpClient.cs Co-authored-by: Stephen Halter <halter73@gmail.com> * Update src/ModelContextProtocol.AspNetCore/StreamableHttpHandler.cs Co-authored-by: Stephen Halter <halter73@gmail.com> * Make init callback asynchronous. * Address feedback. --------- Co-authored-by: Stephen Halter <halter73@gmail.com>
This commit is contained in:
Eirik Tsarpalis
GitHub
parent
fa017c04c3
commit
d7d9eab3f8
11 changed files with 166 additions and 21 deletions
|
|
@ -45,6 +45,8 @@ internal sealed partial class AutoDetectingClientSessionTransport : ITransport
|
|||
|
||||
public ChannelReader<JsonRpcMessage> MessageReader => _messageChannel.Reader;
|
||||
|
||||
string? ITransport.SessionId => ActiveTransport?.SessionId;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public Task SendMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
|
||||
{
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
using Microsoft.Extensions.Logging;
|
||||
using ModelContextProtocol.Protocol;
|
||||
using System.Diagnostics;
|
||||
using System.Text.Json;
|
||||
|
||||
namespace ModelContextProtocol.Client;
|
||||
|
|
@ -93,6 +94,20 @@ internal sealed partial class McpClient : McpEndpoint, IMcpClient
|
|||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public string? SessionId
|
||||
{
|
||||
get
|
||||
{
|
||||
if (_sessionTransport is null)
|
||||
{
|
||||
throw new InvalidOperationException("Must have already initialized a session when invoking this property.");
|
||||
}
|
||||
|
||||
return _sessionTransport.SessionId;
|
||||
}
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public ServerCapabilities ServerCapabilities => _serverCapabilities ?? throw new InvalidOperationException("The client is not connected.");
|
||||
|
||||
|
|
|
|||
|
|
@ -27,7 +27,6 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
private readonly CancellationTokenSource _connectionCts;
|
||||
private readonly ILogger _logger;
|
||||
|
||||
private string? _mcpSessionId;
|
||||
private Task? _getReceiveTask;
|
||||
|
||||
public StreamableHttpClientSessionTransport(
|
||||
|
|
@ -85,7 +84,7 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
},
|
||||
};
|
||||
|
||||
CopyAdditionalHeaders(httpRequestMessage.Headers, _options.AdditionalHeaders, _mcpSessionId);
|
||||
CopyAdditionalHeaders(httpRequestMessage.Headers, _options.AdditionalHeaders, SessionId);
|
||||
|
||||
var response = await _httpClient.SendAsync(httpRequestMessage, HttpCompletionOption.ResponseHeadersRead, cancellationToken).ConfigureAwait(false);
|
||||
|
||||
|
|
@ -124,7 +123,7 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
// We've successfully initialized! Copy session-id and start GET request if any.
|
||||
if (response.Headers.TryGetValues("mcp-session-id", out var sessionIdValues))
|
||||
{
|
||||
_mcpSessionId = sessionIdValues.FirstOrDefault();
|
||||
SessionId = sessionIdValues.FirstOrDefault();
|
||||
}
|
||||
|
||||
_getReceiveTask = ReceiveUnsolicitedMessagesAsync();
|
||||
|
|
@ -170,7 +169,7 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
|
|||
// Send a GET request to handle any unsolicited messages not sent over a POST response.
|
||||
using var request = new HttpRequestMessage(HttpMethod.Get, _options.Endpoint);
|
||||
request.Headers.Accept.Add(s_textEventStreamMediaType);
|
||||
CopyAdditionalHeaders(request.Headers, _options.AdditionalHeaders, _mcpSessionId);
|
||||
CopyAdditionalHeaders(request.Headers, _options.AdditionalHeaders, SessionId);
|
||||
|
||||
using var response = await _httpClient.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, _connectionCts.Token).ConfigureAwait(false);
|
||||
|
||||
|
|
|
|||
|
|
@ -28,6 +28,14 @@ namespace ModelContextProtocol;
|
|||
/// </remarks>
|
||||
public interface IMcpEndpoint : IAsyncDisposable
|
||||
{
|
||||
/// <summary>Gets an identifier associated with the current MCP session.</summary>
|
||||
/// <remarks>
|
||||
/// Typically populated in transports supporting multiple sessions such as Streamable HTTP or SSE.
|
||||
/// Can return <see langword="null"/> if the session hasn't initialized or if the transport doesn't
|
||||
/// support multiple sessions (as is the case with STDIO).
|
||||
/// </remarks>
|
||||
string? SessionId { get; }
|
||||
|
||||
/// <summary>
|
||||
/// Sends a JSON-RPC request to the connected endpoint and waits for a response.
|
||||
/// </summary>
|
||||
|
|
|
|||
|
|
@ -25,6 +25,14 @@ namespace ModelContextProtocol.Protocol;
|
|||
/// </remarks>
|
||||
public interface ITransport : IAsyncDisposable
|
||||
{
|
||||
/// <summary>Gets an identifier associated with the current MCP session.</summary>
|
||||
/// <remarks>
|
||||
/// Typically populated in transports supporting multiple sessions such as Streamable HTTP or SSE.
|
||||
/// Can return <see langword="null"/> if the session hasn't initialized or if the transport doesn't
|
||||
/// support multiple sessions (as is the case with STDIO).
|
||||
/// </remarks>
|
||||
string? SessionId { get; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets a channel reader for receiving messages from the transport.
|
||||
/// </summary>
|
||||
|
|
|
|||
|
|
@ -59,6 +59,9 @@ public abstract partial class TransportBase : ITransport
|
|||
/// <summary>Gets the logger used by this transport.</summary>
|
||||
private protected ILogger Logger => _logger;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public virtual string? SessionId { get; protected set; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets the name that identifies this transport endpoint in logs.
|
||||
/// </summary>
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ namespace ModelContextProtocol.Server;
|
|||
internal sealed class DestinationBoundMcpServer(McpServer server, ITransport? transport) : IMcpServer
|
||||
{
|
||||
public string EndpointName => server.EndpointName;
|
||||
public string? SessionId => transport?.SessionId ?? server.SessionId;
|
||||
public ClientCapabilities? ClientCapabilities => server.ClientCapabilities;
|
||||
public Implementation? ClientInfo => server.ClientInfo;
|
||||
public McpServerOptions ServerOptions => server.ServerOptions;
|
||||
|
|
|
|||
|
|
@ -96,6 +96,10 @@ internal sealed class McpServer : McpEndpoint, IMcpServer
|
|||
InitializeSession(transport);
|
||||
}
|
||||
|
||||
/// <inheritdoc/>
|
||||
public string? SessionId => _sessionTransport.SessionId;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public ServerCapabilities ServerCapabilities { get; } = new();
|
||||
|
||||
/// <inheritdoc />
|
||||
|
|
|
|||
|
|
@ -23,7 +23,8 @@ namespace ModelContextProtocol.Server;
|
|||
/// These messages should be passed to <see cref="OnMessageReceivedAsync(JsonRpcMessage, CancellationToken)"/>.
|
||||
/// Defaults to "/message".
|
||||
/// </param>
|
||||
public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? messageEndpoint = "/message") : ITransport
|
||||
/// <param name="sessionId">The identifier corresponding to the current MCP session.</param>
|
||||
public sealed class SseResponseStreamTransport(Stream sseResponseStream, string? messageEndpoint = "/message", string? sessionId = null) : ITransport
|
||||
{
|
||||
private readonly SseWriter _sseWriter = new(messageEndpoint);
|
||||
private readonly Channel<JsonRpcMessage> _incomingChannel = Channel.CreateBounded<JsonRpcMessage>(new BoundedChannelOptions(1)
|
||||
|
|
@ -49,6 +50,9 @@ public sealed class SseResponseStreamTransport(Stream sseResponseStream, string?
|
|||
/// <inheritdoc/>
|
||||
public ChannelReader<JsonRpcMessage> MessageReader => _incomingChannel.Reader;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public string? SessionId { get; } = sessionId;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public async ValueTask DisposeAsync()
|
||||
{
|
||||
|
|
|
|||
|
|
@ -18,6 +18,8 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
|
|||
|
||||
public ChannelReader<JsonRpcMessage> MessageReader => throw new NotSupportedException("JsonRpcMessage.RelatedTransport should only be used for sending messages.");
|
||||
|
||||
string? ITransport.SessionId => parentTransport.SessionId;
|
||||
|
||||
/// <returns>
|
||||
/// True, if data was written to the respond body.
|
||||
/// False, if nothing was written because the request body did not contain any <see cref="JsonRpcRequest"/> messages to respond to.
|
||||
|
|
@ -79,10 +81,11 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
|
|||
{
|
||||
_pendingRequest = request.Id;
|
||||
|
||||
// Store client capabilities so they can be serialized by "stateless" callers for use in later requests.
|
||||
if (parentTransport.Stateless && request.Method == RequestMethods.Initialize)
|
||||
// Invoke the initialize request callback if applicable.
|
||||
if (parentTransport.OnInitRequestReceived is { } onInitRequest && request.Method == RequestMethods.Initialize)
|
||||
{
|
||||
parentTransport.InitializeRequest = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.JsonContext.Default.InitializeRequestParams);
|
||||
var initializeRequest = JsonSerializer.Deserialize(request.Params, McpJsonUtilities.JsonContext.Default.InitializeRequestParams);
|
||||
await onInitRequest(initializeRequest).ConfigureAwait(false);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -46,15 +46,18 @@ public sealed class StreamableHttpServerTransport : ITransport
|
|||
public bool Stateless { get; init; }
|
||||
|
||||
/// <summary>
|
||||
/// Gets the initialize request if it was received by <see cref="HandlePostRequest(IDuplexPipe, CancellationToken)"/> and <see cref="Stateless"/> is set to <see langword="true"/>.
|
||||
/// Gets or sets a callback to be invoked before handling the initialize request.
|
||||
/// </summary>
|
||||
public InitializeRequestParams? InitializeRequest { get; internal set; }
|
||||
public Func<InitializeRequestParams?, ValueTask>? OnInitRequestReceived { get; set; }
|
||||
|
||||
/// <inheritdoc/>
|
||||
public ChannelReader<JsonRpcMessage> MessageReader => _incomingChannel.Reader;
|
||||
|
||||
internal ChannelWriter<JsonRpcMessage> MessageWriter => _incomingChannel.Writer;
|
||||
|
||||
/// <inheritdoc/>
|
||||
public string? SessionId { get; set; }
|
||||
|
||||
/// <summary>
|
||||
/// Handles an optional SSE GET request a client using the Streamable HTTP transport might make by
|
||||
/// writing any unsolicited JSON-RPC messages sent via <see cref="SendMessageAsync"/>
|
||||
|
|
|
|||
Loading…
Reference in a new issue