Amend test fix

Author: MackinnonBuck

tests/ModelContextProtocol.AspNetCore.Tests/ResumabilityIntegrationTests.cs

@@ -344,8 +344,8 @@ public async Task Client_CanResumeUnsolicitedMessageStream_AfterDisconnection()
             return default;
         });
 
-        // Wait for the client's GET SSE stream to be established before sending notifications
-        await faultingStreamHandler.WaitForStreamAsync(TestContext.Current.CancellationToken);
+        // Wait for the client's unsolicited message stream to be established before sending notifications
+        await faultingStreamHandler.WaitForUnsolicitedMessageStreamAsync(TestContext.Current.CancellationToken);
 
         // Send a custom notification to the client on the unsolicited message stream
         await server.SendNotificationAsync(CustomNotificationMethod, new JsonObject { ["message"] = InitialMessage }, cancellationToken: TestContext.Current.CancellationToken);

tests/ModelContextProtocol.AspNetCore.Tests/Utils/FaultingStreamHandler.cs

@@ -11,10 +11,12 @@ internal sealed class FaultingStreamHandler : DelegatingHandler
 {
     private FaultingStream? _lastStream;
     private TaskCompletionSource? _reconnectTcs;
-    private TaskCompletionSource _streamAvailableTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    private TaskCompletionSource _unsolicitedMessageStreamReadyTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
-    public Task WaitForStreamAsync(CancellationToken cancellationToken = default)
-        => _streamAvailableTcs.Task.WaitAsync(cancellationToken);
+    public Task WaitForUnsolicitedMessageStreamAsync(CancellationToken cancellationToken = default)
+        => _unsolicitedMessageStreamReadyTcs.Task.WaitAsync(cancellationToken);
+
+    internal void SignalUnsolicitedMessageStreamReady() => _unsolicitedMessageStreamReadyTcs.TrySetResult();
 
     public async Task<ReconnectAttempt> TriggerFaultAsync(CancellationToken cancellationToken)
     {
@@ -28,7 +30,9 @@ public async Task<ReconnectAttempt> TriggerFaultAsync(CancellationToken cancella
             throw new InvalidOperationException("Cannot trigger a fault while already waiting for reconnection.");
         }
 
-        _streamAvailableTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+        // Reset the TCS so we can wait for the reconnected unsolicited message stream
+        _unsolicitedMessageStreamReadyTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+
         _reconnectTcs = new();
         await _lastStream.TriggerFaultAsync(cancellationToken);
 
@@ -51,6 +55,7 @@ protected override async Task<HttpResponseMessage> SendAsync(
             _reconnectTcs = null;
         }
 
+        var isGetRequest = request.Method == HttpMethod.Get;
         var response = await base.SendAsync(request, cancellationToken);
 
         // Only wrap SSE streams (text/event-stream)
@@ -69,7 +74,12 @@ protected override async Task<HttpResponseMessage> SendAsync(
 
             response.Content = newContent;
 
-            _streamAvailableTcs.TrySetResult();
+            // For GET requests (unsolicited message stream), set up the stream to signal
+            // when first data is read. This ensures the server's transport handler is ready.
+            if (isGetRequest)
+            {
+                _lastStream.SetReadyCallback(SignalUnsolicitedMessageStreamReady);
+            }
         }
 
         return response;
@@ -96,10 +106,14 @@ private sealed class FaultingStream(Stream innerStream) : Stream
     {
         private readonly CancellationTokenSource _cts = new();
         private TaskCompletionSource? _faultTcs;
+        private Action? _readyCallback;
+        private bool _readySignaled;
         private bool _disposed;
 
         public bool IsDisposed => _disposed;
 
+        public void SetReadyCallback(Action callback) => _readyCallback = callback;
+
         public async Task TriggerFaultAsync(CancellationToken cancellationToken)
         {
             if (_faultTcs is not null)
@@ -138,6 +152,12 @@ public override async ValueTask<int> ReadAsync(Memory<byte> buffer, Cancellation
 
                 _cts.Token.ThrowIfCancellationRequested();
 
+                if (bytesRead > 0 && !_readySignaled)
+                {
+                    _readySignaled = true;
+                    _readyCallback?.Invoke();
+                }
+
                 return bytesRead;
             }
             catch (OperationCanceledException) when (_cts.IsCancellationRequested)