diff --git a/ModelContextProtocol.slnx b/ModelContextProtocol.slnx index 9020d2fbe..b679d22c4 100644 --- a/ModelContextProtocol.slnx +++ b/ModelContextProtocol.slnx @@ -40,6 +40,7 @@ + @@ -73,6 +74,7 @@ + diff --git a/samples/AspNetCoreMcpClient/AspNetCoreMcpClient.csproj b/samples/AspNetCoreMcpClient/AspNetCoreMcpClient.csproj new file mode 100644 index 000000000..7d7a8102a --- /dev/null +++ b/samples/AspNetCoreMcpClient/AspNetCoreMcpClient.csproj @@ -0,0 +1,13 @@ + + + + net10.0 + enable + enable + + + + + + + diff --git a/samples/AspNetCoreMcpClient/ElicitationBroker.cs b/samples/AspNetCoreMcpClient/ElicitationBroker.cs new file mode 100644 index 000000000..7d9485ac2 --- /dev/null +++ b/samples/AspNetCoreMcpClient/ElicitationBroker.cs @@ -0,0 +1,108 @@ +using ModelContextProtocol.Protocol; + +namespace AspNetCoreMcpClient; + +/// +/// Bridges an MCP elicitation request to an application-owned HTTP interaction. +/// +public sealed class ElicitationBroker +{ + private readonly Dictionary _pending = new(StringComparer.Ordinal); + private readonly Lock _lock = new(); + + public async ValueTask RequestAsync( + string sessionId, + ElicitRequestParams? request, + CancellationToken cancellationToken) + { + var pending = new PendingRequest(Guid.NewGuid(), request); + + lock (_lock) + { + if (!_pending.TryAdd(sessionId, pending)) + { + throw new InvalidOperationException("Only one elicitation can be pending for an application session."); + } + } + + try + { + return await pending.Completion.Task.WaitAsync(cancellationToken).ConfigureAwait(false); + } + finally + { + lock (_lock) + { + if (_pending.TryGetValue(sessionId, out var current) && ReferenceEquals(current, pending)) + { + _pending.Remove(sessionId); + } + } + } + } + + public PendingElicitation? GetPending(string sessionId) + { + lock (_lock) + { + return _pending.TryGetValue(sessionId, out var pending) + ? new PendingElicitation(pending.Id, pending.Request) + : null; + } + } + + public bool TryRespond(string sessionId, Guid requestId, ElicitResult response) + { + ArgumentNullException.ThrowIfNull(response); + + lock (_lock) + { + return _pending.TryGetValue(sessionId, out var pending) && + pending.Id == requestId && + pending.Completion.TrySetResult(response); + } + } + + public void Cancel(string sessionId) + { + lock (_lock) + { + if (_pending.Remove(sessionId, out var pending)) + { + pending.Completion.TrySetCanceled(); + } + } + } + + /// + /// Cancels every pending elicitation. Call this during application shutdown so that MCP operations blocked on an + /// unanswered elicitation unblock and the session registry can dispose its clients. + /// + public int CancelAll() + { + lock (_lock) + { + var pending = _pending.Values.ToArray(); + _pending.Clear(); + + foreach (var request in pending) + { + request.Completion.TrySetCanceled(); + } + + return pending.Length; + } + } + + private sealed class PendingRequest(Guid id, ElicitRequestParams? request) + { + public Guid Id { get; } = id; + + public ElicitRequestParams? Request { get; } = request; + + public TaskCompletionSource Completion { get; } = + new(TaskCreationOptions.RunContinuationsAsynchronously); + } +} + +public sealed record PendingElicitation(Guid Id, ElicitRequestParams? Request); diff --git a/samples/AspNetCoreMcpClient/InlineProgress.cs b/samples/AspNetCoreMcpClient/InlineProgress.cs new file mode 100644 index 000000000..356714da7 --- /dev/null +++ b/samples/AspNetCoreMcpClient/InlineProgress.cs @@ -0,0 +1,6 @@ +namespace AspNetCoreMcpClient; + +public sealed class InlineProgress(Action report) : IProgress +{ + public void Report(T value) => report(value); +} diff --git a/samples/AspNetCoreMcpClient/McpClientCleanupService.cs b/samples/AspNetCoreMcpClient/McpClientCleanupService.cs new file mode 100644 index 000000000..e01cffcd5 --- /dev/null +++ b/samples/AspNetCoreMcpClient/McpClientCleanupService.cs @@ -0,0 +1,22 @@ +namespace AspNetCoreMcpClient; + +public sealed class McpClientCleanupService( + SessionClientRegistry registry, + IConfiguration configuration, + ILogger logger) : BackgroundService +{ + protected override async Task ExecuteAsync(CancellationToken stoppingToken) + { + var interval = TimeSpan.FromMinutes(configuration.GetValue("McpServer:CleanupIntervalMinutes", 1)); + using var timer = new PeriodicTimer(interval); + + while (await timer.WaitForNextTickAsync(stoppingToken).ConfigureAwait(false)) + { + var removed = await registry.RemoveIdleAsync().ConfigureAwait(false); + if (removed > 0) + { + logger.LogInformation("Disposed {Count} idle MCP client sessions.", removed); + } + } + } +} diff --git a/samples/AspNetCoreMcpClient/McpClientConnection.cs b/samples/AspNetCoreMcpClient/McpClientConnection.cs new file mode 100644 index 000000000..3073bf7ed --- /dev/null +++ b/samples/AspNetCoreMcpClient/McpClientConnection.cs @@ -0,0 +1,20 @@ +using ModelContextProtocol.Client; + +namespace AspNetCoreMcpClient; + +public sealed class McpClientConnection(McpClient client, HttpClientTransport transport) : IAsyncDisposable +{ + public McpClient Client { get; } = client; + + public async ValueTask DisposeAsync() + { + try + { + await Client.DisposeAsync().ConfigureAwait(false); + } + finally + { + await transport.DisposeAsync().ConfigureAwait(false); + } + } +} diff --git a/samples/AspNetCoreMcpClient/Program.cs b/samples/AspNetCoreMcpClient/Program.cs new file mode 100644 index 000000000..00774d697 --- /dev/null +++ b/samples/AspNetCoreMcpClient/Program.cs @@ -0,0 +1,160 @@ +using AspNetCoreMcpClient; +using ModelContextProtocol; +using ModelContextProtocol.Client; +using ModelContextProtocol.Protocol; +using System.Collections.Concurrent; +using System.Text.Json; + +var builder = WebApplication.CreateBuilder(args); + +var endpoint = new Uri(builder.Configuration["McpServer:Endpoint"] ?? "http://localhost:3001"); +var idleTimeout = TimeSpan.FromMinutes(builder.Configuration.GetValue("McpServer:IdleTimeoutMinutes", 20)); + +builder.Services.AddHttpClient("mcp-server"); +builder.Services.AddSingleton(); +builder.Services.AddSingleton(serviceProvider => +{ + var httpClientFactory = serviceProvider.GetRequiredService(); + var loggerFactory = serviceProvider.GetRequiredService(); + var elicitationBroker = serviceProvider.GetRequiredService(); + + return new SessionClientRegistry( + async (sessionId, cancellationToken) => + { + var httpClient = httpClientFactory.CreateClient("mcp-server"); + var transport = new HttpClientTransport( + new() + { + Endpoint = endpoint, + Name = $"ASP.NET Core session {sessionId}", + TransportMode = HttpTransportMode.StreamableHttp, + }, + httpClient, + loggerFactory, + ownsHttpClient: true); + + try + { + var client = await McpClient.CreateAsync( + transport, + new() + { + ClientInfo = new() { Name = "AspNetCoreMcpClient", Version = "1.0.0" }, + Handlers = new() + { + ElicitationHandler = (request, token) => + elicitationBroker.RequestAsync(sessionId, request, token), + }, + }, + loggerFactory, + cancellationToken); + + return new McpClientConnection(client, transport); + } + catch + { + await transport.DisposeAsync(); + throw; + } + }, + TimeProvider.System, + idleTimeout); +}); +builder.Services.AddHostedService(); + +var app = builder.Build(); + +// Disposing the registry waits for each session's in-flight operation to finish. A tool call blocked on an unanswered +// elicitation would never finish, so release those requests before shutdown disposes the singleton registry. +app.Lifetime.ApplicationStopping.Register(() => +{ + var broker = app.Services.GetRequiredService(); + var canceled = broker.CancelAll(); + if (canceled > 0) + { + app.Logger.LogInformation("Canceled {Count} pending elicitations during shutdown.", canceled); + } +}); + +app.MapGet("/tools", async ( + HttpContext context, + SessionClientRegistry registry, + CancellationToken cancellationToken) => +{ + var sessionId = GetDemoSessionId(context); + var tools = await registry.ExecuteAsync( + sessionId, + async (connection, token) => await connection.Client.ListToolsAsync(cancellationToken: token), + cancellationToken); + + return tools.Select(tool => new { tool.Name, tool.Description }); +}); + +app.MapPost("/tools/{toolName}", async ( + string toolName, + JsonElement? arguments, + HttpContext context, + SessionClientRegistry registry, + CancellationToken cancellationToken) => +{ + var sessionId = GetDemoSessionId(context); + var toolArguments = arguments is { ValueKind: JsonValueKind.Object } + ? arguments.Value.Deserialize>() + : null; + + // Progress notifications arrive on the MCP session's message loop, so use a thread-safe collection. + var progressUpdates = new ConcurrentQueue(); + var progress = new InlineProgress(progressUpdates.Enqueue); + var result = await registry.ExecuteAsync( + sessionId, + async (connection, token) => await connection.Client.CallToolAsync( + toolName, + toolArguments, + progress, + cancellationToken: token), + cancellationToken); + + return Results.Ok(new { Result = result, Progress = progressUpdates }); +}); + +app.MapGet("/elicitation", (HttpContext context, ElicitationBroker broker) => +{ + var pending = broker.GetPending(GetDemoSessionId(context)); + return pending is null ? Results.NoContent() : Results.Ok(pending); +}); + +app.MapPost("/elicitation/{requestId:guid}", ( + Guid requestId, + ElicitResult response, + HttpContext context, + ElicitationBroker broker) => +{ + return broker.TryRespond(GetDemoSessionId(context), requestId, response) + ? Results.Accepted() + : Results.NotFound(); +}); + +app.MapDelete("/session", async ( + HttpContext context, + SessionClientRegistry registry, + ElicitationBroker broker) => +{ + var sessionId = GetDemoSessionId(context); + broker.Cancel(sessionId); + return await registry.RemoveAsync(sessionId) ? Results.NoContent() : Results.NotFound(); +}); + +app.Run(); + +static string GetDemoSessionId(HttpContext context) +{ + const string HeaderName = "X-Demo-User"; + var sessionId = context.Request.Headers[HeaderName].ToString(); + if (string.IsNullOrWhiteSpace(sessionId) || sessionId.Length > 128) + { + throw new BadHttpRequestException($"Provide a non-empty {HeaderName} header of at most 128 characters."); + } + + // A production application should derive this key from authenticated server-side identity/session state. + return sessionId; +} diff --git a/samples/AspNetCoreMcpClient/README.md b/samples/AspNetCoreMcpClient/README.md new file mode 100644 index 000000000..786b0809a --- /dev/null +++ b/samples/AspNetCoreMcpClient/README.md @@ -0,0 +1,59 @@ +# ASP.NET Core MCP client with per-user sessions + +This sample shows an ASP.NET Core Web API acting as an MCP client while keeping one MCP connection per application user session. It is intended for applications that need session continuity for server-to-client features such as elicitation and progress notifications. + +The sample deliberately separates the **application session** from the MCP protocol session. `SessionClientRegistry` owns the mapping and provides: + +- lazy, single initialization of a client for each application session; +- serialization of operations within one session while allowing different sessions to run concurrently; +- explicit session removal and deterministic asynchronous disposal; +- automatic removal of idle sessions; and +- cleanup of every remaining client during application shutdown. + +## Run the sample + +Start an HTTP MCP server, such as `AspNetCoreMcpServer`, then run this project: + +```bash +dotnet run --project samples/AspNetCoreMcpServer +dotnet run --project samples/AspNetCoreMcpClient +``` + +The default MCP endpoint is `http://localhost:3001`. Change `McpServer:Endpoint` in `appsettings.json` when needed. + +The HTTP examples use `X-Demo-User` solely to make session reuse visible without adding an authentication system: + +```bash +curl -H "X-Demo-User: alice" http://localhost:5000/tools + +curl -X POST \ + -H "Content-Type: application/json" \ + -H "X-Demo-User: alice" \ + -d '{"message":"hello"}' \ + http://localhost:5000/tools/echo + +curl -X DELETE -H "X-Demo-User: alice" http://localhost:5000/session +``` + +Use the URL printed by `dotnet run` if it differs from port 5000. + +## Elicitation flow + +When the MCP server sends an elicitation request during a tool call, `ElicitationBroker` holds that request while the application's frontend collects an answer: + +1. The frontend polls `GET /elicitation` with the same application-session identity. +2. A `200` response contains the pending request and its ID; `204` means there is no pending request. +3. The frontend posts an `ElicitResult` to `POST /elicitation/{requestId}`. +4. The original tool call resumes and returns its response. + +The broker intentionally permits one pending elicitation per application session because the registry serializes that session's MCP operations. + +Because the registry waits for a session's in-flight operation before disposing its client, a tool call blocked on an unanswered elicitation would otherwise stall `DELETE /session` and application shutdown. `DELETE /session` cancels the session's pending elicitation first, and the sample registers an `ApplicationStopping` callback that cancels all pending elicitations before the registry is disposed. + +## Production considerations + +- **Never trust a caller-provided session header.** Replace `X-Demo-User` with a key derived from authenticated, server-side identity or session state. Do not use access tokens or other secrets as dictionary keys. +- The registry is in-memory and therefore single-node. For multiple application instances, use sticky routing so one user's requests reach the owning process, or implement distributed ownership and session resumption. A distributed cache alone cannot store a live `McpClient` connection. +- Choose an idle timeout that fits both application behavior and upstream resource limits. Explicitly remove the session at logout when possible. +- Operations are serialized per user to protect application-level session state. If your use case permits concurrent MCP requests, adjust the registry policy rather than creating duplicate clients. +- The sample collects progress updates for a compact response. A real frontend would normally stream them with Server-Sent Events, WebSockets, or another application channel. diff --git a/samples/AspNetCoreMcpClient/SessionClientRegistry.cs b/samples/AspNetCoreMcpClient/SessionClientRegistry.cs new file mode 100644 index 000000000..097c1999c --- /dev/null +++ b/samples/AspNetCoreMcpClient/SessionClientRegistry.cs @@ -0,0 +1,196 @@ +using System.Collections.Concurrent; + +namespace AspNetCoreMcpClient; + +/// +/// Owns one asynchronously disposable client per application session and serializes operations within each session. +/// +/// The client type stored for each session. +public sealed class SessionClientRegistry : IAsyncDisposable + where TClient : IAsyncDisposable +{ + private readonly ConcurrentDictionary _entries = new(StringComparer.Ordinal); + private readonly Func> _clientFactory; + private readonly TimeProvider _timeProvider; + private readonly TimeSpan _idleTimeout; + private int _disposed; + + public SessionClientRegistry( + Func> clientFactory, + TimeProvider timeProvider, + TimeSpan idleTimeout) + { + ArgumentNullException.ThrowIfNull(clientFactory); + ArgumentNullException.ThrowIfNull(timeProvider); + ArgumentOutOfRangeException.ThrowIfLessThanOrEqual(idleTimeout, TimeSpan.Zero); + + _clientFactory = clientFactory; + _timeProvider = timeProvider; + _idleTimeout = idleTimeout; + } + + /// + /// Runs an operation against the session's client. Operations for different sessions can run concurrently, while + /// operations for the same session are serialized. + /// + public async Task ExecuteAsync( + string sessionId, + Func> operation, + CancellationToken cancellationToken = default) + { + ArgumentException.ThrowIfNullOrWhiteSpace(sessionId); + ArgumentNullException.ThrowIfNull(operation); + ObjectDisposedException.ThrowIf(Volatile.Read(ref _disposed) != 0, this); + + while (true) + { + var entry = _entries.GetOrAdd(sessionId, _ => new Entry(_timeProvider.GetUtcNow())); + await entry.Gate.WaitAsync(cancellationToken).ConfigureAwait(false); + + try + { + if (entry.Removed) + { + continue; + } + + if (Volatile.Read(ref _disposed) != 0) + { + RemoveExact(sessionId, entry); + entry.Removed = true; + throw new ObjectDisposedException(GetType().FullName); + } + + if (entry.Client is null) + { + try + { + entry.Client = await _clientFactory(sessionId, cancellationToken).ConfigureAwait(false); + } + catch + { + RemoveExact(sessionId, entry); + entry.Removed = true; + throw; + } + } + + entry.LastAccess = _timeProvider.GetUtcNow(); + try + { + return await operation(entry.Client, cancellationToken).ConfigureAwait(false); + } + finally + { + entry.LastAccess = _timeProvider.GetUtcNow(); + } + } + finally + { + entry.Gate.Release(); + } + } + } + + /// Removes and disposes a session client, if one exists. + public async Task RemoveAsync(string sessionId) + { + ArgumentException.ThrowIfNullOrWhiteSpace(sessionId); + + if (!_entries.TryGetValue(sessionId, out var entry) || !RemoveExact(sessionId, entry)) + { + return false; + } + + await DisposeEntryAsync(entry).ConfigureAwait(false); + return true; + } + + /// Removes clients that have not been used within the configured idle timeout. + public async Task RemoveIdleAsync() + { + var removed = 0; + var now = _timeProvider.GetUtcNow(); + + foreach (var pair in _entries) + { + var entry = pair.Value; + if (!await entry.Gate.WaitAsync(0).ConfigureAwait(false)) + { + continue; + } + + try + { + if (!entry.Removed && now - entry.LastAccess >= _idleTimeout && RemoveExact(pair.Key, entry)) + { + entry.Removed = true; + if (entry.Client is not null) + { + await entry.Client.DisposeAsync().ConfigureAwait(false); + entry.Client = default; + } + + removed++; + } + } + finally + { + entry.Gate.Release(); + } + } + + return removed; + } + + public async ValueTask DisposeAsync() + { + if (Interlocked.Exchange(ref _disposed, 1) != 0) + { + return; + } + + while (!_entries.IsEmpty) + { + foreach (var pair in _entries) + { + if (RemoveExact(pair.Key, pair.Value)) + { + await DisposeEntryAsync(pair.Value).ConfigureAwait(false); + } + } + } + } + + private bool RemoveExact(string sessionId, Entry entry) => + ((ICollection>)_entries).Remove(new(sessionId, entry)); + + private static async Task DisposeEntryAsync(Entry entry) + { + await entry.Gate.WaitAsync().ConfigureAwait(false); + try + { + entry.Removed = true; + if (entry.Client is not null) + { + await entry.Client.DisposeAsync().ConfigureAwait(false); + entry.Client = default; + } + } + finally + { + entry.Gate.Release(); + } + } + + private sealed class Entry(DateTimeOffset lastAccess) + { + public SemaphoreSlim Gate { get; } = new(1, 1); + + public TClient? Client { get; set; } + + public DateTimeOffset LastAccess { get; set; } = lastAccess; + + public bool Removed { get; set; } + } +} diff --git a/samples/AspNetCoreMcpClient/appsettings.json b/samples/AspNetCoreMcpClient/appsettings.json new file mode 100644 index 000000000..bf302d586 --- /dev/null +++ b/samples/AspNetCoreMcpClient/appsettings.json @@ -0,0 +1,7 @@ +{ + "McpServer": { + "Endpoint": "http://localhost:3001", + "IdleTimeoutMinutes": 20, + "CleanupIntervalMinutes": 1 + } +} diff --git a/tests/AspNetCoreMcpClient.Tests/AspNetCoreMcpClient.Tests.csproj b/tests/AspNetCoreMcpClient.Tests/AspNetCoreMcpClient.Tests.csproj new file mode 100644 index 000000000..92e9b8205 --- /dev/null +++ b/tests/AspNetCoreMcpClient.Tests/AspNetCoreMcpClient.Tests.csproj @@ -0,0 +1,25 @@ + + + + net10.0 + enable + enable + false + true + + + + + + + + runtime; build; native; contentfiles; analyzers; buildtransitive + all + + + + + + + + diff --git a/tests/AspNetCoreMcpClient.Tests/ElicitationBrokerTests.cs b/tests/AspNetCoreMcpClient.Tests/ElicitationBrokerTests.cs new file mode 100644 index 000000000..41dc8cc73 --- /dev/null +++ b/tests/AspNetCoreMcpClient.Tests/ElicitationBrokerTests.cs @@ -0,0 +1,52 @@ +using AspNetCoreMcpClient; +using ModelContextProtocol.Protocol; + +namespace AspNetCoreMcpClient.Tests; + +public class ElicitationBrokerTests +{ + [Fact] + public async Task Response_CompletesMatchingPendingRequest() + { + var broker = new ElicitationBroker(); + var requestTask = broker.RequestAsync("alice", new() { Message = "Choose" }, TestContext.Current.CancellationToken).AsTask(); + var pending = broker.GetPending("alice"); + + Assert.NotNull(pending); + var expected = new ElicitResult { Action = "accept" }; + Assert.True(broker.TryRespond("alice", pending.Id, expected)); + + Assert.Same(expected, await requestTask); + Assert.Null(broker.GetPending("alice")); + } + + [Fact] + public async Task Cancellation_RemovesPendingRequest() + { + var broker = new ElicitationBroker(); + using var cts = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + var requestTask = broker.RequestAsync("alice", new() { Message = "Choose" }, cts.Token).AsTask(); + + Assert.NotNull(broker.GetPending("alice")); + cts.Cancel(); + + await Assert.ThrowsAnyAsync(() => requestTask); + Assert.Null(broker.GetPending("alice")); + } + + [Fact] + public async Task CancelAll_ReleasesEveryPendingRequest() + { + var broker = new ElicitationBroker(); + var alice = broker.RequestAsync("alice", new() { Message = "Choose" }, TestContext.Current.CancellationToken).AsTask(); + var bob = broker.RequestAsync("bob", new() { Message = "Choose" }, TestContext.Current.CancellationToken).AsTask(); + + Assert.Equal(2, broker.CancelAll()); + + await Assert.ThrowsAnyAsync(() => alice); + await Assert.ThrowsAnyAsync(() => bob); + Assert.Null(broker.GetPending("alice")); + Assert.Null(broker.GetPending("bob")); + Assert.Equal(0, broker.CancelAll()); + } +} diff --git a/tests/AspNetCoreMcpClient.Tests/GlobalUsings.cs b/tests/AspNetCoreMcpClient.Tests/GlobalUsings.cs new file mode 100644 index 000000000..c802f4480 --- /dev/null +++ b/tests/AspNetCoreMcpClient.Tests/GlobalUsings.cs @@ -0,0 +1 @@ +global using Xunit; diff --git a/tests/AspNetCoreMcpClient.Tests/SessionClientRegistryTests.cs b/tests/AspNetCoreMcpClient.Tests/SessionClientRegistryTests.cs new file mode 100644 index 000000000..b38f9709a --- /dev/null +++ b/tests/AspNetCoreMcpClient.Tests/SessionClientRegistryTests.cs @@ -0,0 +1,111 @@ +using AspNetCoreMcpClient; +using Microsoft.Extensions.Time.Testing; + +namespace AspNetCoreMcpClient.Tests; + +public class SessionClientRegistryTests +{ + [Fact] + public async Task ConcurrentOperationsForOneSession_CreateOneClientAndDoNotOverlap() + { + var cancellationToken = TestContext.Current.CancellationToken; + var factoryCalls = 0; + var concurrentOperations = 0; + var maximumConcurrency = 0; + var releaseFirstOperation = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var firstOperationStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + await using var registry = new SessionClientRegistry( + (_, _) => Task.FromResult(new FakeClient(Interlocked.Increment(ref factoryCalls))), + TimeProvider.System, + TimeSpan.FromMinutes(5)); + + async ValueTask Operation(FakeClient client, CancellationToken _) + { + var current = Interlocked.Increment(ref concurrentOperations); + maximumConcurrency = Math.Max(maximumConcurrency, current); + firstOperationStarted.TrySetResult(); + await releaseFirstOperation.Task.WaitAsync(cancellationToken); + Interlocked.Decrement(ref concurrentOperations); + return client.Id; + } + + var first = registry.ExecuteAsync("alice", Operation, cancellationToken); + await firstOperationStarted.Task.WaitAsync(cancellationToken); + var second = registry.ExecuteAsync("alice", Operation, cancellationToken); + releaseFirstOperation.SetResult(); + + var results = await Task.WhenAll(first, second); + + Assert.Equal([1, 1], results); + Assert.Equal(1, factoryCalls); + Assert.Equal(1, maximumConcurrency); + } + + [Fact] + public async Task DifferentSessions_CanRunConcurrently() + { + var cancellationToken = TestContext.Current.CancellationToken; + var concurrentOperations = 0; + var bothStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + + await using var registry = new SessionClientRegistry( + (_, _) => Task.FromResult(new FakeClient(1)), + TimeProvider.System, + TimeSpan.FromMinutes(5)); + + async ValueTask Operation(FakeClient _, CancellationToken cancellationToken) + { + if (Interlocked.Increment(ref concurrentOperations) == 2) + { + bothStarted.SetResult(); + } + + await bothStarted.Task.WaitAsync(cancellationToken); + Interlocked.Decrement(ref concurrentOperations); + return 1; + } + + await Task.WhenAll( + registry.ExecuteAsync("alice", Operation, cancellationToken), + registry.ExecuteAsync("bob", Operation, cancellationToken)); + } + + [Fact] + public async Task RemoveAndIdleCleanup_DisposeClients() + { + var cancellationToken = TestContext.Current.CancellationToken; + var timeProvider = new FakeTimeProvider(); + var clients = new List(); + await using var registry = new SessionClientRegistry( + (_, _) => + { + var client = new FakeClient(clients.Count + 1); + clients.Add(client); + return Task.FromResult(client); + }, + timeProvider, + TimeSpan.FromMinutes(5)); + + await registry.ExecuteAsync("explicit", static (client, _) => ValueTask.FromResult(client.Id), cancellationToken); + await registry.ExecuteAsync("idle", static (client, _) => ValueTask.FromResult(client.Id), cancellationToken); + + Assert.True(await registry.RemoveAsync("explicit")); + timeProvider.Advance(TimeSpan.FromMinutes(6)); + Assert.Equal(1, await registry.RemoveIdleAsync()); + Assert.All(clients, client => Assert.True(client.IsDisposed)); + } + + private sealed class FakeClient(int id) : IAsyncDisposable + { + public int Id { get; } = id; + + public bool IsDisposed { get; private set; } + + public ValueTask DisposeAsync() + { + IsDisposed = true; + return default; + } + } +}