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

tests/ModelContextProtocol.AspNetCore.Tests/MapMcpTests.cs

@@ -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}";
+        }
+    }
 }