Don't conflate request and response IDs in Streamable HTTP transports (#475)

This commit is contained in:
Stephen Halter 2025-06-03 12:17:11 -07:00 • committed by GitHub
commit 046903c27b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 147 additions and 26 deletions

View file

@ -96,17 +96,17 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
}
var rpcRequest = message as JsonRpcRequest;
JsonRpcMessage? rpcResponseCandidate = null;
JsonRpcMessageWithId? rpcResponseOrError = null;
if (response.Content.Headers.ContentType?.MediaType == "application/json")
{
var responseContent = await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false);
rpcResponseCandidate = await ProcessMessageAsync(responseContent, cancellationToken).ConfigureAwait(false);
rpcResponseOrError = await ProcessMessageAsync(responseContent, rpcRequest, cancellationToken).ConfigureAwait(false);
}
else if (response.Content.Headers.ContentType?.MediaType == "text/event-stream")
{
using var responseBodyStream = await response.Content.ReadAsStreamAsync(cancellationToken);
rpcResponseCandidate = await ProcessSseResponseAsync(responseBodyStream, rpcRequest, cancellationToken).ConfigureAwait(false);
rpcResponseOrError = await ProcessSseResponseAsync(responseBodyStream, rpcRequest, cancellationToken).ConfigureAwait(false);
}
if (rpcRequest is null)
@ -114,12 +114,12 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
return response;
}
if (rpcResponseCandidate is not JsonRpcMessageWithId messageWithId || messageWithId.Id != rpcRequest.Id)
if (rpcResponseOrError is null)
{
throw new McpException($"Streamable HTTP POST response completed without a reply to request with ID: {rpcRequest.Id}");
}
if (rpcRequest.Method == RequestMethods.Initialize && rpcResponseCandidate is JsonRpcResponse)
if (rpcRequest.Method == RequestMethods.Initialize && rpcResponseOrError is JsonRpcResponse)
{
// We've successfully initialized! Copy session-id and start GET request if any.
if (response.Headers.TryGetValues("mcp-session-id", out var sessionIdValues))
@ -193,20 +193,20 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
continue;
}
var message = await ProcessMessageAsync(sseEvent.Data, cancellationToken).ConfigureAwait(false);
var rpcResponseOrError = await ProcessMessageAsync(sseEvent.Data, relatedRpcRequest, cancellationToken).ConfigureAwait(false);
// The server SHOULD end the response here anyway, but we won't leave it to chance. This transport makes
// The server SHOULD end the HTTP response body here anyway, but we won't leave it to chance. This transport makes
// a GET request for any notifications that might need to be sent after the completion of each POST.
if (message is JsonRpcMessageWithId messageWithId && relatedRpcRequest?.Id == messageWithId.Id)
if (rpcResponseOrError is not null)
{
return messageWithId;
return rpcResponseOrError;
}
}
return null;
}
private async Task<JsonRpcMessage?> ProcessMessageAsync(string data, CancellationToken cancellationToken)
private async Task<JsonRpcMessageWithId?> ProcessMessageAsync(string data, JsonRpcRequest? relatedRpcRequest, CancellationToken cancellationToken)
{
try
{
@ -218,7 +218,12 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
}
await WriteMessageAsync(message, cancellationToken).ConfigureAwait(false);
return message;
if (message is JsonRpcResponse or JsonRpcError &&
message is JsonRpcMessageWithId rpcResponseOrError &&
rpcResponseOrError.Id == relatedRpcRequest?.Id)
{
return rpcResponseOrError;
}
}
catch (JsonException ex)
{

View file

@ -60,7 +60,7 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
{
yield return message;
if (message.Data is JsonRpcMessageWithId response && response.Id == _pendingRequest)
if (message.Data is JsonRpcResponse or JsonRpcError && ((JsonRpcMessageWithId)message.Data).Id == _pendingRequest)
{
// Complete the SSE response stream now that all pending requests have been processed.
break;

View file

@ -12,16 +12,16 @@ public class LoggedTest : IDisposable
{
CurrentTestOutputHelper = testOutputHelper,
};
LoggerProvider = new XunitLoggerProvider(_delegatingTestOutputHelper);
XunitLoggerProvider = new XunitLoggerProvider(_delegatingTestOutputHelper);
LoggerFactory = Microsoft.Extensions.Logging.LoggerFactory.Create(builder =>
{
builder.AddProvider(LoggerProvider);
builder.AddProvider(XunitLoggerProvider);
});
}
public ITestOutputHelper TestOutputHelper => _delegatingTestOutputHelper;
public ILoggerFactory LoggerFactory { get; }
public ILoggerProvider LoggerProvider { get; }
public ILoggerFactory LoggerFactory { get; set; }
public ILoggerProvider XunitLoggerProvider { get; }
public virtual void Dispose()
{

View file

@ -52,7 +52,7 @@ public class MapMcpSseTests(ITestOutputHelper outputHelper) : MapMcpTests(output
await app.StartAsync(TestContext.Current.CancellationToken);
var mcpClient = await ConnectAsync(requestPath);
await using var mcpClient = await ConnectAsync(requestPath);
Assert.Equal("TestCustomRouteServer", mcpClient.ServerInfo.Name);
}

View file

@ -31,7 +31,7 @@ public class MapMcpStreamableHttpTests(ITestOutputHelper outputHelper) : MapMcpT
await app.StartAsync(TestContext.Current.CancellationToken);
var mcpClient = await ConnectAsync(requestPath);
await using var mcpClient = await ConnectAsync(requestPath);
Assert.Equal("TestCustomRouteServer", mcpClient.ServerInfo.Name);
}
@ -135,7 +135,7 @@ public class MapMcpStreamableHttpTests(ITestOutputHelper outputHelper) : MapMcpT
await app.StartAsync(TestContext.Current.CancellationToken);
await using var mcpClient = await ConnectAsync(options: new()
await using var mcpClient = await ConnectAsync(transportOptions: new()
{
Endpoint = new Uri("http://localhost/sse"),
TransportMode = HttpTransportMode.Sse

View file

@ -1,9 +1,12 @@
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.AspNetCore.Tests.Utils;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;
using ModelContextProtocol.Tests.Utils;
using System.ComponentModel;
using System.Net;
using System.Security.Claims;
@ -20,18 +23,21 @@ public abstract class MapMcpTests(ITestOutputHelper testOutputHelper) : KestrelI
options.Stateless = Stateless;
}
protected async Task<IMcpClient> ConnectAsync(string? path = null, SseClientTransportOptions? options = null)
protected async Task<IMcpClient> ConnectAsync(
string? path = null,
SseClientTransportOptions? transportOptions = null,
McpClientOptions? clientOptions = null)
{
// Default behavior when no options are provided
path ??= UseStreamableHttp ? "/" : "/sse";
await using var transport = new SseClientTransport(options ?? new SseClientTransportOptions()
await using var transport = new SseClientTransport(transportOptions ?? new SseClientTransportOptions()
{
Endpoint = new Uri($"http://localhost{path}"),
TransportMode = UseStreamableHttp ? HttpTransportMode.StreamableHttp : HttpTransportMode.Sse,
}, HttpClient, LoggerFactory);
return await McpClientFactory.CreateAsync(transport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken);
return await McpClientFactory.CreateAsync(transport, clientOptions, LoggerFactory, TestContext.Current.CancellationToken);
}
[Fact]
@ -71,7 +77,7 @@ public abstract class MapMcpTests(ITestOutputHelper testOutputHelper) : KestrelI
await app.StartAsync(TestContext.Current.CancellationToken);
var mcpClient = await ConnectAsync();
await using var mcpClient = await ConnectAsync();
var response = await mcpClient.CallToolAsync(
"EchoWithUserName",
@ -111,13 +117,90 @@ public abstract class MapMcpTests(ITestOutputHelper testOutputHelper) : KestrelI
Assert.Equal(HttpStatusCode.Forbidden, httpRequestException.StatusCode);
}
protected ClaimsPrincipal CreateUser(string name)
[Fact]
public async Task Sampling_DoesNotCloseStream_Prematurely()
{
Assert.SkipWhen(Stateless, "Sampling is not supported in stateless mode.");
Builder.Services.AddMcpServer().WithHttpTransport(ConfigureStateless).WithTools<SamplingRegressionTools>();
var mockLoggerProvider = new MockLoggerProvider();
Builder.Logging.AddProvider(mockLoggerProvider);
Builder.Logging.SetMinimumLevel(LogLevel.Debug);
await using var app = Builder.Build();
// Reset the LoggerFactory used by the client to use the MockLoggerProvider as well.
LoggerFactory = app.Services.GetRequiredService<ILoggerFactory>();
app.MapMcp();
await app.StartAsync(TestContext.Current.CancellationToken);
var sampleCount = 0;
var clientOptions = new McpClientOptions
{
Capabilities = new()
{
Sampling = new()
{
SamplingHandler = async (parameters, _, _) =>
{
Assert.NotNull(parameters?.Messages);
var message = Assert.Single(parameters.Messages);
Assert.Equal(Role.User, message.Role);
Assert.Equal("text", message.Content.Type);
Assert.Equal("Test prompt for sampling", message.Content.Text);
sampleCount++;
return new CreateMessageResult
{
Model = "test-model",
Role = Role.Assistant,
Content = new Content
{
Type = "text",
Text = "Sampling response from client"
}
};
},
},
},
};
await using var mcpClient = await ConnectAsync(clientOptions: clientOptions);
var result = await mcpClient.CallToolAsync("sampling-tool", new Dictionary<string, object?>
{
["prompt"] = "Test prompt for sampling"
}, cancellationToken: TestContext.Current.CancellationToken);
Assert.NotNull(result);
Assert.False(result.IsError);
var textContent = Assert.Single(result.Content);
Assert.Equal("text", textContent.Type);
Assert.Equal("Sampling completed successfully. Client responded: Sampling response from client", textContent.Text);
Assert.Equal(2, sampleCount);
// Verify that the tool call and the sampling request both used the same ID to ensure we cover against regressions.
// https://github.com/modelcontextprotocol/csharp-sdk/issues/464
Assert.Single(mockLoggerProvider.LogMessages, m =>
m.Category == "ModelContextProtocol.Client.McpClient" &&
m.Message.Contains("request '2' for method 'tools/call'"));
Assert.Single(mockLoggerProvider.LogMessages, m =>
m.Category == "ModelContextProtocol.Server.McpServer" &&
m.Message.Contains("request '2' for method 'sampling/createMessage'"));
}
private ClaimsPrincipal CreateUser(string name)
=> new ClaimsPrincipal(new ClaimsIdentity(
[new Claim("name", name), new Claim(ClaimTypes.NameIdentifier, name)],
"TestAuthType", "name", "role"));
[McpServerToolType]
protected class EchoHttpContextUserTools(IHttpContextAccessor contextAccessor)
private class EchoHttpContextUserTools(IHttpContextAccessor contextAccessor)
{
[McpServerTool, Description("Echoes the input back to the client with their user name.")]
public string EchoWithUserName(string message)
@ -127,4 +210,37 @@ public abstract class MapMcpTests(ITestOutputHelper testOutputHelper) : KestrelI
return $"{userName}: {message}";
}
}
[McpServerToolType]
private class SamplingRegressionTools
{
[McpServerTool(Name = "sampling-tool")]
public static async Task<string> SamplingToolAsync(IMcpServer server, string prompt, CancellationToken cancellationToken)
{
// This tool reproduces the scenario described in https://github.com/modelcontextprotocol/csharp-sdk/issues/464
// 1. The client calls tool with request ID 2, because it's the first request after the initialize request.
// 2. This tool makes two sampling requests which use IDs 1 and 2.
// 3. In the old buggy Streamable HTTP transport code, this would close the SSE response stream,
// because the second sampling request used an ID matching the tool call.
var samplingRequest = new CreateMessageRequestParams
{
Messages = [
new SamplingMessage
{
Role = Role.User,
Content = new Content
{
Type = "text",
Text = prompt
},
}
],
};
await server.SampleAsync(samplingRequest, cancellationToken);
var samplingResult = await server.SampleAsync(samplingRequest, cancellationToken);
return $"Sampling completed successfully. Client responded: {samplingResult.Content.Text}";
}
}
}

View file

@ -18,7 +18,7 @@ public class KestrelInMemoryTest : LoggedTest
Builder = WebApplication.CreateSlimBuilder();
Builder.Services.RemoveAll<IConnectionListenerFactory>();
Builder.Services.AddSingleton<IConnectionListenerFactory>(_inMemoryTransport);
Builder.Services.AddSingleton(LoggerProvider);
Builder.Services.AddSingleton(XunitLoggerProvider);
HttpClient = new HttpClient(new SocketsHttpHandler()
{