Remove Channel

Author: halter73

src/ModelContextProtocol.Core/Server/McpServerImpl.cs

@@ -7,7 +7,6 @@
 using System.Text.Json;
 using System.Text.Json.Nodes;
 using System.Text.Json.Serialization.Metadata;
-using System.Threading.Channels;
 
 namespace ModelContextProtocol.Server;
 
@@ -211,10 +210,7 @@ public override async ValueTask DisposeAsync()
         {
             if (_mrtrContinuations.TryRemove(kvp.Key, out var continuation))
             {
-                foreach (var exchange in continuation.PendingExchanges)
-                {
-                    exchange.ResponseTcs.TrySetCanceled();
-                }
+                continuation.PendingExchange.ResponseTcs.TrySetCanceled();
             }
         }
 
@@ -1192,19 +1188,20 @@ private void WrapHandlerWithMrtr(string method)
                     inputResponses = JsonSerializer.Deserialize(responsesNode, McpJsonUtilities.JsonContext.Default.IDictionaryStringInputResponse);
                 }
 
-                // Complete pending exchanges with the client's responses.
-                foreach (var exchange in continuation.PendingExchanges)
+                // Prepare for the next potential exchange before resuming the handler.
+                continuation.MrtrContext.ResetForNextExchange();
+
+                // Complete the pending exchange with the client's response.
+                var exchange = continuation.PendingExchange;
+                if (inputResponses is not null &&
+                    inputResponses.TryGetValue(exchange.Key, out var response))
                 {
-                    if (inputResponses is not null &&
-                        inputResponses.TryGetValue(exchange.Key, out var response))
-                    {
-                        exchange.ResponseTcs.TrySetResult(response);
-                    }
-                    else
-                    {
-                        exchange.ResponseTcs.TrySetException(
-                            new McpProtocolException($"Missing input response for key '{exchange.Key}'.", McpErrorCode.InvalidParams));
-                    }
+                    exchange.ResponseTcs.TrySetResult(response);
+                }
+                else
+                {
+                    exchange.ResponseTcs.TrySetException(
+                        new McpProtocolException($"Missing input response for key '{exchange.Key}'.", McpErrorCode.InvalidParams));
                 }
 
                 // Race again: handler completion vs new exchange.
@@ -1228,7 +1225,7 @@ private void WrapHandlerWithMrtr(string method)
             Task<JsonNode?> handlerTask;
             try
             {
-                handlerTask = InvokeOriginalHandlerAsync(originalHandler, request, mrtrContext, cancellationToken);
+                handlerTask = originalHandler(request, cancellationToken);
             }
             finally
             {
@@ -1241,31 +1238,7 @@ private void WrapHandlerWithMrtr(string method)
     }
 
     /// <summary>
-    /// Invokes the original request handler and marks the MrtrContext as complete when done.
-    /// </summary>
-    private static async Task<JsonNode?> InvokeOriginalHandlerAsync(
-        Func<JsonRpcRequest, CancellationToken, Task<JsonNode?>> handler,
-        JsonRpcRequest request,
-        MrtrContext mrtrContext,
-        CancellationToken cancellationToken)
-    {
-        try
-        {
-            return await handler(request, cancellationToken).ConfigureAwait(false);
-        }
-        catch (Exception ex)
-        {
-            mrtrContext.Fault(ex);
-            throw;
-        }
-        finally
-        {
-            mrtrContext.Complete();
-        }
-    }
-
-    /// <summary>
-    /// Races between handler completion and the MrtrContext exchange channel.
+    /// Races between handler completion and the MrtrContext exchange TCS.
     /// If the handler completes, returns its result. If an exchange arrives (handler needs input),
     /// builds and returns an IncompleteResult and stores the continuation for future retries.
     /// </summary>
@@ -1280,10 +1253,7 @@ private void WrapHandlerWithMrtr(string method)
             return await handlerTask.ConfigureAwait(false);
         }
 
-        // Start reading from the exchange channel.
-        var readTask = mrtrContext.ExchangeReader.ReadAsync(cancellationToken).AsTask();
-
-        var completedTask = await Task.WhenAny(handlerTask, readTask).ConfigureAwait(false);
+        var completedTask = await Task.WhenAny(handlerTask, mrtrContext.ExchangeTask).ConfigureAwait(false);
 
         if (completedTask == handlerTask)
         {
@@ -1292,40 +1262,17 @@ private void WrapHandlerWithMrtr(string method)
         }
 
         // Exchange arrived - handler needs input from the client.
-        MrtrExchange firstExchange;
-        try
-        {
-            firstExchange = await readTask.ConfigureAwait(false);
-        }
-        catch (ChannelClosedException)
-        {
-            // Channel was closed (handler completed between WhenAny and ReadAsync).
-            return await handlerTask.ConfigureAwait(false);
-        }
-
-        // Collect all currently available exchanges (handles concurrent ElicitAsync/SampleAsync calls).
-        var exchanges = new List<MrtrExchange> { firstExchange };
-        while (mrtrContext.ExchangeReader.TryRead(out var additionalExchange))
-        {
-            exchanges.Add(additionalExchange);
-        }
-
-        // Build the IncompleteResult with input requests.
-        var inputRequests = new Dictionary<string, InputRequest>(exchanges.Count);
-        foreach (var exchange in exchanges)
-        {
-            inputRequests[exchange.Key] = exchange.InputRequest;
-        }
+        var exchange = await mrtrContext.ExchangeTask.ConfigureAwait(false);
 
         var correlationId = Guid.NewGuid().ToString("N");
         var incompleteResult = new IncompleteResult
         {
-            InputRequests = inputRequests,
+            InputRequests = new Dictionary<string, InputRequest> { [exchange.Key] = exchange.InputRequest },
             RequestState = correlationId,
         };
 
         // Store the continuation so the retry can resume the handler.
-        _mrtrContinuations[correlationId] = new MrtrContinuation(handlerTask, mrtrContext, exchanges);
+        _mrtrContinuations[correlationId] = new MrtrContinuation(handlerTask, mrtrContext, exchange);
 
         return JsonSerializer.SerializeToNode(incompleteResult, McpJsonUtilities.JsonContext.Default.IncompleteResult);
     }

src/ModelContextProtocol.Core/Server/MrtrContext.cs

@@ -1,5 +1,4 @@
 using System.Text.Json.Nodes;
-using System.Threading.Channels;
 using ModelContextProtocol.Protocol;
 
 namespace ModelContextProtocol.Server;
@@ -8,8 +7,9 @@ namespace ModelContextProtocol.Server;
 /// Manages the MRTR (Multi Round-Trip Request) coordination between a handler and the pipeline.
 /// When a handler calls <see cref="McpServer.ElicitAsync(ModelContextProtocol.Protocol.ElicitRequestParams, System.Threading.CancellationToken)"/> or
 /// <see cref="McpServer.SampleAsync(ModelContextProtocol.Protocol.CreateMessageRequestParams, System.Threading.CancellationToken)"/>,
-/// the handler writes to the channel and suspends on a TCS. The pipeline reads from the channel,
-/// sends an <see cref="IncompleteResult"/>, and later completes the TCS when the retry arrives.
+/// the handler sets the exchange TCS and suspends on a response TCS. The pipeline detects the exchange
+/// via <see cref="ExchangeTask"/>, sends an <see cref="IncompleteResult"/>, and later completes the
+/// response TCS when the retry arrives.
 /// </summary>
 internal sealed class MrtrContext
 {
@@ -18,15 +18,14 @@ internal sealed class MrtrContext
     /// </summary>
     internal const string ExperimentalCapabilityKey = "mrtr";
 
-    private readonly Channel<MrtrExchange> _exchanges = Channel.CreateUnbounded<MrtrExchange>(
-        new UnboundedChannelOptions { SingleReader = true });
+    private TaskCompletionSource<MrtrExchange> _exchangeTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
     private int _nextInputRequestId;
 
     /// <summary>
-    /// Gets the channel reader for consuming exchanges produced by the handler.
+    /// Gets a task that completes when the handler produces an exchange (calls ElicitAsync/SampleAsync/RequestRootsAsync).
     /// </summary>
-    public ChannelReader<MrtrExchange> ExchangeReader => _exchanges.Reader;
+    public Task<MrtrExchange> ExchangeTask => _exchangeTcs.Task;
 
     /// <summary>
     /// Called by <see cref="McpServer.ElicitAsync(ModelContextProtocol.Protocol.ElicitRequestParams, System.Threading.CancellationToken)"/>
@@ -36,26 +35,30 @@ internal sealed class MrtrContext
     /// <param name="inputRequest">The input request describing what the server needs.</param>
     /// <param name="cancellationToken">A token to cancel the wait for input.</param>
     /// <returns>The client's response to the input request.</returns>
+    /// <exception cref="InvalidOperationException">A concurrent server-to-client request is already pending.</exception>
     public async Task<InputResponse> RequestInputAsync(InputRequest inputRequest, CancellationToken cancellationToken)
     {
-        var key = $"input_{Interlocked.Increment(ref _nextInputRequestId)}";
+        var tcs = _exchangeTcs;
+        if (tcs.Task.IsCompleted)
+        {
+            throw new InvalidOperationException("Concurrent server-to-client requests are not supported. Await each ElicitAsync, SampleAsync, or RequestRootsAsync call before making another.");
+        }
 
+        var key = $"input_{Interlocked.Increment(ref _nextInputRequestId)}";
         var exchange = new MrtrExchange(key, inputRequest);
-
-        await _exchanges.Writer.WriteAsync(exchange, cancellationToken).ConfigureAwait(false);
+        tcs.TrySetResult(exchange);
 
         return await exchange.ResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
     }
 
     /// <summary>
-    /// Signals that the handler has completed normally.
-    /// </summary>
-    public void Complete() => _exchanges.Writer.TryComplete();
-
-    /// <summary>
-    /// Signals that the handler has faulted.
+    /// Prepares the context for the next round of exchange after a retry arrives.
+    /// Must be called before completing the previous exchange's response TCS.
     /// </summary>
-    public void Fault(Exception exception) => _exchanges.Writer.TryComplete(exception);
+    public void ResetForNextExchange()
+    {
+        _exchangeTcs = new TaskCompletionSource<MrtrExchange>(TaskCreationOptions.RunContinuationsAsynchronously);
+    }
 }
 
 /// <summary>
@@ -93,11 +96,11 @@ public MrtrExchange(string key, InputRequest inputRequest)
 /// </summary>
 internal sealed class MrtrContinuation
 {
-    public MrtrContinuation(Task<JsonNode?> handlerTask, MrtrContext mrtrContext, IReadOnlyList<MrtrExchange> pendingExchanges)
+    public MrtrContinuation(Task<JsonNode?> handlerTask, MrtrContext mrtrContext, MrtrExchange pendingExchange)
     {
         HandlerTask = handlerTask;
         MrtrContext = mrtrContext;
-        PendingExchanges = pendingExchanges;
+        PendingExchange = pendingExchange;
     }
 
     /// <summary>
@@ -111,7 +114,7 @@ public MrtrContinuation(Task<JsonNode?> handlerTask, MrtrContext mrtrContext, IR
     public MrtrContext MrtrContext { get; }
 
     /// <summary>
-    /// The exchanges that are awaiting responses from the client.
+    /// The exchange that is awaiting a response from the client.
     /// </summary>
-    public IReadOnlyList<MrtrExchange> PendingExchanges { get; }
+    public MrtrExchange PendingExchange { get; }
 }