Add tests to verify McpClient.DisposeAsync doesn't hang (#1252)
This commit is contained in:
parent
27b5eb8031
commit
b4c7ded398
2 changed files with 224 additions and 0 deletions
|
|
@ -412,4 +412,139 @@ public class MapMcpStreamableHttpTests(ITestOutputHelper outputHelper) : MapMcpT
|
|||
Assert.True(wasPostRequest, "POST request was not made");
|
||||
Assert.True(wasDeleteRequest, "DELETE request was not made");
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task DisposeAsync_DoesNotHang_WhenOwnsSessionIsFalse()
|
||||
{
|
||||
Assert.SkipWhen(Stateless, "Stateless mode doesn't support session management.");
|
||||
|
||||
var getResponseStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
|
||||
Builder.Services.AddMcpServer().WithHttpTransport(ConfigureStateless).WithTools<ClaimsPrincipalTools>();
|
||||
|
||||
await using var app = Builder.Build();
|
||||
|
||||
// Track when the GET SSE response starts being written, which indicates
|
||||
// the server's HandleGetRequestAsync has fully initialized the SSE writer.
|
||||
app.Use(next =>
|
||||
{
|
||||
return async context =>
|
||||
{
|
||||
if (context.Request.Method == HttpMethods.Get)
|
||||
{
|
||||
context.Response.OnStarting(() =>
|
||||
{
|
||||
getResponseStarted.TrySetResult();
|
||||
return Task.CompletedTask;
|
||||
});
|
||||
}
|
||||
await next(context);
|
||||
};
|
||||
});
|
||||
|
||||
app.MapMcp();
|
||||
await app.StartAsync(TestContext.Current.CancellationToken);
|
||||
|
||||
await using var transport = new HttpClientTransport(new()
|
||||
{
|
||||
Endpoint = new("http://localhost:5000/"),
|
||||
TransportMode = HttpTransportMode.StreamableHttp,
|
||||
OwnsSession = false,
|
||||
}, HttpClient, LoggerFactory);
|
||||
|
||||
var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
// Call a tool to ensure the session is fully established
|
||||
var result = await client.CallToolAsync(
|
||||
"echo_claims_principal",
|
||||
new Dictionary<string, object?>() { ["message"] = "Hello!" },
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
Assert.NotNull(result);
|
||||
|
||||
// Wait for the GET SSE stream to be fully established on the server
|
||||
await getResponseStarted.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
|
||||
|
||||
// This should not hang. The issue reports that DisposeAsync hangs indefinitely
|
||||
// when OwnsSession is false. Use a timeout to detect the hang.
|
||||
await client.DisposeAsync().AsTask().WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task DisposeAsync_DoesNotHang_WhenOwnsSessionIsFalse_WithUnsolicitedMessages()
|
||||
{
|
||||
Assert.SkipWhen(Stateless, "Stateless mode doesn't support session management.");
|
||||
|
||||
var getResponseStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
var serverTcs = new TaskCompletionSource<McpServer>(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
|
||||
Builder.Services.AddMcpServer().WithHttpTransport(opts =>
|
||||
{
|
||||
ConfigureStateless(opts);
|
||||
opts.RunSessionHandler = async (context, server, cancellationToken) =>
|
||||
{
|
||||
serverTcs.TrySetResult(server);
|
||||
await server.RunAsync(cancellationToken);
|
||||
};
|
||||
}).WithTools<ClaimsPrincipalTools>();
|
||||
|
||||
await using var app = Builder.Build();
|
||||
|
||||
// Track when the GET SSE response starts being written, which indicates
|
||||
// the server's HandleGetRequestAsync has fully initialized the SSE writer.
|
||||
app.Use(next =>
|
||||
{
|
||||
return async context =>
|
||||
{
|
||||
if (context.Request.Method == HttpMethods.Get)
|
||||
{
|
||||
context.Response.OnStarting(() =>
|
||||
{
|
||||
getResponseStarted.TrySetResult();
|
||||
return Task.CompletedTask;
|
||||
});
|
||||
}
|
||||
await next(context);
|
||||
};
|
||||
});
|
||||
|
||||
app.MapMcp();
|
||||
await app.StartAsync(TestContext.Current.CancellationToken);
|
||||
|
||||
await using var transport = new HttpClientTransport(new()
|
||||
{
|
||||
Endpoint = new("http://localhost:5000/"),
|
||||
TransportMode = HttpTransportMode.StreamableHttp,
|
||||
OwnsSession = false,
|
||||
}, HttpClient, LoggerFactory);
|
||||
|
||||
var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
||||
var result = await client.CallToolAsync(
|
||||
"echo_claims_principal",
|
||||
new Dictionary<string, object?>() { ["message"] = "Hello!" },
|
||||
cancellationToken: TestContext.Current.CancellationToken);
|
||||
Assert.NotNull(result);
|
||||
|
||||
// Wait for the GET SSE stream to be fully established on the server
|
||||
await getResponseStarted.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
|
||||
|
||||
// Register a handler on the client to detect when the notification is received
|
||||
var notificationReceived = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
await using var handlerRegistration = client.RegisterNotificationHandler("notifications/tools/list_changed", (notification, ct) =>
|
||||
{
|
||||
notificationReceived.TrySetResult();
|
||||
return default;
|
||||
});
|
||||
|
||||
// Get the server instance and send an unsolicited notification by modifying tools
|
||||
var server = await serverTcs.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
|
||||
await server.SendNotificationAsync("notifications/tools/list_changed", TestContext.Current.CancellationToken);
|
||||
|
||||
// Wait for the client to actually receive the notification
|
||||
await notificationReceived.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
|
||||
|
||||
// Dispose should still not hang
|
||||
await client.DisposeAsync().AsTask().WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -281,6 +281,95 @@ public class StreamableHttpClientConformanceTests(ITestOutputHelper outputHelper
|
|||
Assert.Contains(nameof(McpClient.ResumeSessionAsync), exception.Message);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public async Task DisposeAsync_DoesNotHang_WhenOwnsSessionIsFalse_WithActiveGetStream()
|
||||
{
|
||||
var getRequestReceived = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
|
||||
|
||||
Builder.Services.Configure<JsonOptions>(options =>
|
||||
{
|
||||
options.SerializerOptions.TypeInfoResolverChain.Add(McpJsonUtilities.DefaultOptions.TypeInfoResolver!);
|
||||
});
|
||||
_app = Builder.Build();
|
||||
|
||||
var echoTool = McpServerTool.Create(Echo, new() { Services = _app.Services });
|
||||
|
||||
_app.MapPost("/mcp", (JsonRpcMessage message, HttpContext context) =>
|
||||
{
|
||||
if (message is not JsonRpcRequest request)
|
||||
{
|
||||
return Results.Accepted();
|
||||
}
|
||||
|
||||
context.Response.Headers.Append("mcp-session-id", "hang-test-session");
|
||||
|
||||
if (request.Method == "initialize")
|
||||
{
|
||||
return Results.Json(new JsonRpcResponse
|
||||
{
|
||||
Id = request.Id,
|
||||
Result = JsonSerializer.SerializeToNode(new InitializeResult
|
||||
{
|
||||
ProtocolVersion = "2024-11-05",
|
||||
Capabilities = new() { Tools = new() },
|
||||
ServerInfo = new Implementation { Name = "hang-test", Version = "0.0.1" },
|
||||
}, McpJsonUtilities.DefaultOptions)
|
||||
});
|
||||
}
|
||||
|
||||
if (request.Method == "tools/list")
|
||||
{
|
||||
return Results.Json(new JsonRpcResponse
|
||||
{
|
||||
Id = request.Id,
|
||||
Result = JsonSerializer.SerializeToNode(new ListToolsResult
|
||||
{
|
||||
Tools = [echoTool.ProtocolTool]
|
||||
}, McpJsonUtilities.DefaultOptions),
|
||||
});
|
||||
}
|
||||
|
||||
return Results.Accepted();
|
||||
});
|
||||
|
||||
// GET handler that keeps the SSE stream open indefinitely (like a real MCP server)
|
||||
_app.MapGet("/mcp", async context =>
|
||||
{
|
||||
context.Response.Headers.ContentType = "text/event-stream";
|
||||
getRequestReceived.TrySetResult();
|
||||
await context.Response.Body.FlushAsync(TestContext.Current.CancellationToken);
|
||||
|
||||
try
|
||||
{
|
||||
await Task.Delay(Timeout.Infinite, context.RequestAborted);
|
||||
}
|
||||
catch (OperationCanceledException)
|
||||
{
|
||||
}
|
||||
});
|
||||
|
||||
await _app.StartAsync(TestContext.Current.CancellationToken);
|
||||
|
||||
await using var transport = new HttpClientTransport(new()
|
||||
{
|
||||
Endpoint = new("http://localhost:5000/mcp"),
|
||||
TransportMode = HttpTransportMode.StreamableHttp,
|
||||
OwnsSession = false,
|
||||
}, HttpClient, LoggerFactory);
|
||||
|
||||
await using (var client = await McpClient.CreateAsync(transport, loggerFactory: LoggerFactory, cancellationToken: TestContext.Current.CancellationToken))
|
||||
{
|
||||
var tools = await client.ListToolsAsync(cancellationToken: TestContext.Current.CancellationToken);
|
||||
Assert.Single(tools);
|
||||
|
||||
// Wait for the GET SSE stream to be established on the server
|
||||
await getRequestReceived.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
|
||||
|
||||
// Dispose should not hang even though the GET stream is actively open
|
||||
await client.DisposeAsync().AsTask().WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken);
|
||||
}
|
||||
}
|
||||
|
||||
private static async Task CallEchoAndValidateAsync(McpClientTool echoTool)
|
||||
{
|
||||
var response = await echoTool.CallAsync(new Dictionary<string, object?>() { ["message"] = "Hello world!" }, cancellationToken: TestContext.Current.CancellationToken);
|
||||
|
|
|
|||
Loading…
Reference in a new issue