PR feedback

MackinnonBuck · MackinnonBuck · commit 5502522c2658 · 2026-01-14T09:21:37.000-08:00
diff --git a/src/Common/ServerSentEvents/ArrayBuffer.cs b/src/Common/ServerSentEvents/ArrayBuffer.cs
@@ -1,7 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-// Copied from https://github.com/dotnet/runtime/blob/main/src/libraries/Common/src/System/Net/ArrayBuffer.cs
+// Copied from https://github.com/dotnet/runtime/blob/dcbf3413c5f7ae431a68fd0d3f09af095b525887/src/libraries/Common/src/System/Net/ArrayBuffer.cs
 
 using System.Buffers;
 using System.Diagnostics;
@@ -26,7 +26,7 @@ internal struct ArrayBuffer : IDisposable
 #if NET
     private static int ArrayMaxLength => Array.MaxLength;
 #else
-        private const int ArrayMaxLength = 0X7FFFFFC7;
+    private const int ArrayMaxLength = 0X7FFFFFC7;
 #endif
 
     private readonly bool _usePool;
diff --git a/src/Common/ServerSentEvents/PooledByteBufferWriter.cs b/src/Common/ServerSentEvents/PooledByteBufferWriter.cs
@@ -1,7 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-// Copied from https://github.com/dotnet/runtime/blob/main/src/libraries/System.Net.ServerSentEvents/src/System/Net/ServerSentEvents/PooledByteBufferWriter.cs
+// Copied from https://github.com/dotnet/runtime/blob/dcbf3413c5f7ae431a68fd0d3f09af095b525887/src/libraries/System.Net.ServerSentEvents/src/System/Net/ServerSentEvents/PooledByteBufferWriter.cs
 
 using System.Buffers;
 using System.Diagnostics;
diff --git a/src/Common/ServerSentEvents/SseEventWriter.cs b/src/Common/ServerSentEvents/SseEventWriter.cs
@@ -1,7 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-// Based on https://github.com/dotnet/runtime/blob/main/src/libraries/System.Net.ServerSentEvents/src/System/Net/ServerSentEvents/SseFormatter.cs
+// Based on https://github.com/dotnet/runtime/blob/dcbf3413c5f7ae431a68fd0d3f09af095b525887/src/libraries/System.Net.ServerSentEvents/src/System/Net/ServerSentEvents/SseFormatter.cs
 
 using System.Buffers;
 using System.Diagnostics;
diff --git a/src/Common/ServerSentEvents/SseEventWriterHelpers.cs b/src/Common/ServerSentEvents/SseEventWriterHelpers.cs
@@ -1,7 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-// Copied from https://github.com/dotnet/runtime/blob/main/src/libraries/System.Net.ServerSentEvents/src/System/Net/ServerSentEvents/Helpers.cs
+// Copied from https://github.com/dotnet/runtime/blob/dcbf3413c5f7ae431a68fd0d3f09af095b525887/src/libraries/System.Net.ServerSentEvents/src/System/Net/ServerSentEvents/Helpers.cs
 
 using System.Buffers;
 using System.Diagnostics;
diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs b/src/ModelContextProtocol.Core/Server/StreamableHttpPostTransport.cs
@@ -10,17 +10,18 @@ namespace ModelContextProtocol.Server;
 /// Handles processing the request/response body pairs for the Streamable HTTP transport.
 /// This is typically used via <see cref="JsonRpcMessageContext.RelatedTransport"/>.
 /// </summary>
-internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, Stream responseStream) : ITransport
+internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, Stream responseStream, CancellationToken sessionCancellationToken) : ITransport
 {
     private readonly SemaphoreSlim _messageLock = new(1, 1);
-    private readonly TaskCompletionSource<bool> _responseTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
-    private readonly SseEventWriter _sseResponseWriter = new(responseStream);
+    private readonly TaskCompletionSource<bool> _httpResponseTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+    private readonly SseEventWriter _httpSseWriter = new(responseStream);
+
+    private TaskCompletionSource<bool>? _storeStreamTcs;
+    private ISseEventStreamWriter? _storeSseWriter;
 
-    private TaskCompletionSource<bool>? _streamTcs;
-    private ISseEventStreamWriter? _sseEventStreamWriter;
     private RequestId _pendingRequest;
     private bool _finalResponseMessageSent;
-    private bool _originalResponseCompleted;
+    private bool _httpResponseCompleted;
 
     public ChannelReader<JsonRpcMessage> MessageReader => throw new NotSupportedException("JsonRpcMessage.Context.RelatedTransport should only be used for sending messages.");
 
@@ -31,7 +32,7 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport
     /// False, if nothing was written because the request body did not contain any <see cref="JsonRpcRequest"/> messages to respond to.
     /// The HTTP application should typically respond with an empty "202 Accepted" response in this scenario.
     /// </returns>
-    public async ValueTask<bool> HandlePostAsync(JsonRpcMessage message, CancellationToken postCancellationToken, CancellationToken sessionCancellationToken)
+    public async ValueTask<bool> HandlePostAsync(JsonRpcMessage message, CancellationToken cancellationToken)
     {
         Debug.Assert(_pendingRequest.Id is null);
 
@@ -57,25 +58,25 @@ public async ValueTask<bool> HandlePostAsync(JsonRpcMessage message, Cancellatio
 
         if (_pendingRequest.Id is null)
         {
-            await parentTransport.MessageWriter.WriteAsync(message, postCancellationToken).ConfigureAwait(false);
+            await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
             return false;
         }
 
-        using (await _messageLock.LockAsync(postCancellationToken).ConfigureAwait(false))
+        using (await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false))
         {
-            var primingItem = await TryStartSseEventStreamAsync(_pendingRequest, sessionCancellationToken).ConfigureAwait(false);
+            var primingItem = await TryStartSseEventStreamAsync(_pendingRequest).ConfigureAwait(false);
             if (primingItem.HasValue)
             {
-                await _sseResponseWriter.WriteAsync(primingItem.Value, postCancellationToken).ConfigureAwait(false);
+                await _httpSseWriter.WriteAsync(primingItem.Value, cancellationToken).ConfigureAwait(false);
             }
 
             // Ensure that we've sent the priming event before processing the incoming request.
-            await parentTransport.MessageWriter.WriteAsync(message, postCancellationToken).ConfigureAwait(false);
+            await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);
         }
 
         // Wait for the response to be written before returning from the handler.
         // This keeps the HTTP response open until the final response message is sent.
-        await _responseTcs.Task.WaitAsync(postCancellationToken).ConfigureAwait(false);
+        await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
 
         return true;
     }
@@ -104,22 +105,22 @@ public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken can
 
             var item = new SseItem<JsonRpcMessage?>(message, SseParser.EventTypeDefault);
 
-            if (_sseEventStreamWriter is not null)
+            if (_storeSseWriter is not null)
             {
-                item = await _sseEventStreamWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);
+                item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);
             }
 
-            if (!_originalResponseCompleted)
+            if (!_httpResponseCompleted)
             {
                 // Only write the message to the response if the response has not completed.
 
                 try
                 {
-                    await _sseResponseWriter.WriteAsync(item, cancellationToken).ConfigureAwait(false);
+                    await _httpSseWriter.WriteAsync(item, cancellationToken).ConfigureAwait(false);
                 }
                 catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
                 {
-                    _responseTcs.TrySetException(ex);
+                    _httpResponseTcs.TrySetException(ex);
                 }
             }
         }
@@ -129,8 +130,8 @@ public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken can
             if ((message is JsonRpcResponse or JsonRpcError) && ((JsonRpcMessageWithId)message).Id == _pendingRequest)
             {
                 _finalResponseMessageSent = true;
-                _responseTcs.TrySetResult(true);
-                _streamTcs?.TrySetResult(true);
+                _httpResponseTcs.TrySetResult(true);
+                _storeStreamTcs?.TrySetResult(true);
             }
         }
     }
@@ -144,61 +145,61 @@ public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationTo
 
         using var _ = await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false);
 
-        if (_sseEventStreamWriter is null)
+        if (_storeSseWriter is null)
         {
             throw new InvalidOperationException($"Polling requires an event stream store to be configured.");
         }
 
         // Send the priming event with the new retry interval.
-        var primingItem = await _sseEventStreamWriter.WriteEventAsync(
+        var primingItem = await _storeSseWriter.WriteEventAsync(
             sseItem: new SseItem<JsonRpcMessage?>() { ReconnectionInterval = retryInterval },
             cancellationToken)
             .ConfigureAwait(false);
 
         // Write to the response stream if it still exists.
-        if (!_originalResponseCompleted)
+        if (!_httpResponseCompleted)
         {
-            await _sseResponseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
+            await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
         }
 
         // Set the mode to 'Polling' so that the replay stream ends as soon as all available messages have been sent.
         // This prevents the client from immediately establishing another long-lived connection.
-        await _sseEventStreamWriter.SetModeAsync(SseEventStreamMode.Polling, cancellationToken).ConfigureAwait(false);
+        await _storeSseWriter.SetModeAsync(SseEventStreamMode.Polling, cancellationToken).ConfigureAwait(false);
 
         // Signal completion so HandlePostAsync can return.
-        _responseTcs.TrySetResult(true);
+        _httpResponseTcs.TrySetResult(true);
     }
 
-    private async ValueTask<SseItem<JsonRpcMessage?>?> TryStartSseEventStreamAsync(RequestId requestId, CancellationToken cancellationToken)
+    private async ValueTask<SseItem<JsonRpcMessage?>?> TryStartSseEventStreamAsync(RequestId requestId)
     {
-        Debug.Assert(_sseEventStreamWriter is null);
+        Debug.Assert(_storeSseWriter is null);
 
-        _sseEventStreamWriter = await parentTransport.TryCreateEventStreamAsync(
+        _storeSseWriter = await parentTransport.TryCreateEventStreamAsync(
             streamId: requestId.Id!.ToString()!,
-            cancellationToken: cancellationToken)
+            cancellationToken: sessionCancellationToken)
             .ConfigureAwait(false);
 
-        if (_sseEventStreamWriter is null)
+        if (_storeSseWriter is null)
         {
             return null;
         }
 
-        _streamTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
-        _ = HandleStreamWriterDisposalAsync(_streamTcs.Task, cancellationToken);
+        _storeStreamTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+        _ = HandleStoreStreamDisposalAsync(_storeStreamTcs.Task);
 
-        return await _sseEventStreamWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);
+        return await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), sessionCancellationToken).ConfigureAwait(false);
 
-        async Task HandleStreamWriterDisposalAsync(Task streamTask, CancellationToken cancellationToken)
+        async Task HandleStoreStreamDisposalAsync(Task streamTask)
         {
             try
             {
-                await streamTask.WaitAsync(cancellationToken).ConfigureAwait(false);
+                await streamTask.WaitAsync(sessionCancellationToken).ConfigureAwait(false);
             }
             finally
             {
                 using var _ = await _messageLock.LockAsync().ConfigureAwait(false);
 
-                await _sseEventStreamWriter!.DisposeAsync().ConfigureAwait(false);
+                await _storeSseWriter!.DisposeAsync().ConfigureAwait(false);
             }
         }
     }
@@ -207,16 +208,16 @@ public async ValueTask DisposeAsync()
     {
         using var _ = await _messageLock.LockAsync().ConfigureAwait(false);
 
-        if (_originalResponseCompleted)
+        if (_httpResponseCompleted)
         {
             return;
         }
 
-        _originalResponseCompleted = true;
+        _httpResponseCompleted = true;
 
-        _responseTcs.TrySetResult(true);
+        _httpResponseTcs.TrySetResult(true);
 
-        _sseResponseWriter.Dispose();
+        _httpSseWriter.Dispose();
 
         // Don't dispose the event stream writer here, as we may continue to write to the event store
         // after disposal if there are pending messages.
diff --git a/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs b/src/ModelContextProtocol.Core/Server/StreamableHttpServerTransport.cs
@@ -36,11 +36,11 @@ public sealed class StreamableHttpServerTransport : ITransport
     private readonly CancellationTokenSource _transportDisposedCts = new();
     private readonly SemaphoreSlim _unsolicitedMessageLock = new(1, 1);
 
-    private SseEventWriter? _sseResponseWriter;
-    private ISseEventStreamWriter? _sseEventStreamWriter;
-    private TaskCompletionSource<bool>? _responseTcs;
-    private bool _getRequestStarted;
-    private bool _originalResponseCompleted;
+    private SseEventWriter? _httpSseWriter;
+    private ISseEventStreamWriter? _storeSseWriter;
+    private TaskCompletionSource<bool>? _httpResponseTcs;
+    private bool _getHttpRequestStarted;
+    private bool _getHttpResponseCompleted;
 
     /// <inheritdoc/>
     public string? SessionId { get; init; }
@@ -111,25 +111,25 @@ public async Task HandleGetRequestAsync(Stream sseResponseStream, CancellationTo
 
         using (await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false))
         {
-            if (_getRequestStarted)
+            if (_getHttpRequestStarted)
             {
                 throw new InvalidOperationException("Session resumption is not yet supported. Please start a new session.");
             }
 
-            _getRequestStarted = true;
-            _sseResponseWriter = new SseEventWriter(sseResponseStream);
-            _responseTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
-            _sseEventStreamWriter = await TryCreateEventStreamAsync(streamId: UnsolicitedMessageStreamId, cancellationToken).ConfigureAwait(false);
-            if (_sseEventStreamWriter is not null)
+            _getHttpRequestStarted = true;
+            _httpSseWriter = new SseEventWriter(sseResponseStream);
+            _httpResponseTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
+            _storeSseWriter = await TryCreateEventStreamAsync(streamId: UnsolicitedMessageStreamId, cancellationToken).ConfigureAwait(false);
+            if (_storeSseWriter is not null)
             {
-                var primingItem = await _sseEventStreamWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);
-                await _sseResponseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
+                var primingItem = await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);
+                await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);
             }
         }
 
         // Wait for the response to be written before returning from the handler.
         // This keeps the HTTP response open until the final response message is sent.
-        await _responseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
+        await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
     }
 
     /// <summary>
@@ -155,14 +155,13 @@ public async Task<bool> HandlePostRequestAsync(JsonRpcMessage message, Stream re
         Throw.IfNull(message);
         Throw.IfNull(responseStream);
 
+        var postTransport = new StreamableHttpPostTransport(this, responseStream, _transportDisposedCts.Token);
         using var postCts = CancellationTokenSource.CreateLinkedTokenSource(_transportDisposedCts.Token, cancellationToken);
-        var postTransport = new StreamableHttpPostTransport(this, responseStream);
         await using (postTransport.ConfigureAwait(false))
         {
             return await postTransport.HandlePostAsync(
                 message,
-                postCancellationToken: postCts.Token,
-                sessionCancellationToken: _transportDisposedCts.Token)
+                cancellationToken: postCts.Token)
                 .ConfigureAwait(false);
         }
     }
@@ -179,34 +178,34 @@ public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken can
 
         using var _ = await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false);
 
-        if (!_getRequestStarted)
+        if (!_getHttpRequestStarted)
         {
             // Clients are not required to make a GET request for unsolicited messages.
             // If no GET request has been made, drop the message.
             return;
         }
 
-        Debug.Assert(_sseResponseWriter is not null);
-        Debug.Assert(_responseTcs is not null);
+        Debug.Assert(_httpSseWriter is not null);
+        Debug.Assert(_httpResponseTcs is not null);
 
         var item = SseItem.Message(message);
 
-        if (_sseEventStreamWriter is not null)
+        if (_storeSseWriter is not null)
         {
-            item = await _sseEventStreamWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);
+            item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);
         }
 
-        if (!_originalResponseCompleted)
+        if (!_getHttpResponseCompleted)
         {
             // Only write the message to the response if the response has not completed.
 
             try
             {
-                await _sseResponseWriter!.WriteAsync(item, cancellationToken).ConfigureAwait(false);
+                await _httpSseWriter!.WriteAsync(item, cancellationToken).ConfigureAwait(false);
             }
             catch (Exception ex) when (!cancellationToken.IsCancellationRequested)
             {
-                _responseTcs!.TrySetException(ex);
+                _httpResponseTcs!.TrySetException(ex);
             }
         }
     }
@@ -216,12 +215,12 @@ public async ValueTask DisposeAsync()
     {
         using var _ = await _unsolicitedMessageLock.LockAsync().ConfigureAwait(false);
 
-        if (_originalResponseCompleted)
+        if (_getHttpResponseCompleted)
         {
             return;
         }
 
-        _originalResponseCompleted = true;
+        _getHttpResponseCompleted = true;
 
         try
         {
@@ -232,12 +231,12 @@ public async ValueTask DisposeAsync()
         {
             try
             {
-                _responseTcs?.TrySetResult(true);
-                _sseResponseWriter?.Dispose();
+                _httpResponseTcs?.TrySetResult(true);
+                _httpSseWriter?.Dispose();
 
-                if (_sseEventStreamWriter is not null)
+                if (_storeSseWriter is not null)
                 {
-                    await _sseEventStreamWriter.DisposeAsync().ConfigureAwait(false);
+                    await _storeSseWriter.DisposeAsync().ConfigureAwait(false);
                 }
             }
             finally

Original file line number	Diff line number	Diff line change
`@@ -10,17 +10,18 @@ namespace ModelContextProtocol.Server;`
`10`	`10`	`/// Handles processing the request/response body pairs for the Streamable HTTP transport.`
`11`	`11`	`/// This is typically used via <see cref="JsonRpcMessageContext.RelatedTransport"/>.`
`12`	`12`	`/// </summary>`
`13`		`-internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, Stream responseStream) : ITransport`
	`13`	`+internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport parentTransport, Stream responseStream, CancellationToken sessionCancellationToken) : ITransport`
`14`	`14`	`{`
`15`	`15`	`private readonly SemaphoreSlim _messageLock = new(1, 1);`
`16`		`- private readonly TaskCompletionSource<bool> _responseTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);`
`17`		`- private readonly SseEventWriter _sseResponseWriter = new(responseStream);`
	`16`	`+ private readonly TaskCompletionSource<bool> _httpResponseTcs = new(TaskCreationOptions.RunContinuationsAsynchronously);`
	`17`	`+ private readonly SseEventWriter _httpSseWriter = new(responseStream);`
	`18`	`+`
	`19`	`+ private TaskCompletionSource<bool>? _storeStreamTcs;`
	`20`	`+ private ISseEventStreamWriter? _storeSseWriter;`
`18`	`21`
`19`		`- private TaskCompletionSource<bool>? _streamTcs;`
`20`		`- private ISseEventStreamWriter? _sseEventStreamWriter;`
`21`	`22`	`private RequestId _pendingRequest;`
`22`	`23`	`private bool _finalResponseMessageSent;`
`23`		`- private bool _originalResponseCompleted;`
	`24`	`+ private bool _httpResponseCompleted;`
`24`	`25`
`25`	`26`	`public ChannelReader<JsonRpcMessage> MessageReader => throw new NotSupportedException("JsonRpcMessage.Context.RelatedTransport should only be used for sending messages.");`
`26`	`27`
`@@ -31,7 +32,7 @@ internal sealed class StreamableHttpPostTransport(StreamableHttpServerTransport`
`31`	`32`	`/// False, if nothing was written because the request body did not contain any <see cref="JsonRpcRequest"/> messages to respond to.`
`32`	`33`	`/// The HTTP application should typically respond with an empty "202 Accepted" response in this scenario.`
`33`	`34`	`/// </returns>`
`34`		`- public async ValueTask<bool> HandlePostAsync(JsonRpcMessage message, CancellationToken postCancellationToken, CancellationToken sessionCancellationToken)`
	`35`	`+ public async ValueTask<bool> HandlePostAsync(JsonRpcMessage message, CancellationToken cancellationToken)`
`35`	`36`	`{`
`36`	`37`	`Debug.Assert(_pendingRequest.Id is null);`
`37`	`38`
`@@ -57,25 +58,25 @@ public async ValueTask<bool> HandlePostAsync(JsonRpcMessage message, Cancellatio`
`57`	`58`
`58`	`59`	`if (_pendingRequest.Id is null)`
`59`	`60`	`{`
`60`		`- await parentTransport.MessageWriter.WriteAsync(message, postCancellationToken).ConfigureAwait(false);`
	`61`	`+ await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);`
`61`	`62`	`return false;`
`62`	`63`	`}`
`63`	`64`
`64`		`- using (await _messageLock.LockAsync(postCancellationToken).ConfigureAwait(false))`
	`65`	`+ using (await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false))`
`65`	`66`	`{`
`66`		`- var primingItem = await TryStartSseEventStreamAsync(_pendingRequest, sessionCancellationToken).ConfigureAwait(false);`
	`67`	`+ var primingItem = await TryStartSseEventStreamAsync(_pendingRequest).ConfigureAwait(false);`
`67`	`68`	`if (primingItem.HasValue)`
`68`	`69`	`{`
`69`		`- await _sseResponseWriter.WriteAsync(primingItem.Value, postCancellationToken).ConfigureAwait(false);`
	`70`	`+ await _httpSseWriter.WriteAsync(primingItem.Value, cancellationToken).ConfigureAwait(false);`
`70`	`71`	`}`
`71`	`72`
`72`	`73`	`// Ensure that we've sent the priming event before processing the incoming request.`
`73`		`- await parentTransport.MessageWriter.WriteAsync(message, postCancellationToken).ConfigureAwait(false);`
	`74`	`+ await parentTransport.MessageWriter.WriteAsync(message, cancellationToken).ConfigureAwait(false);`
`74`	`75`	`}`
`75`	`76`
`76`	`77`	`// Wait for the response to be written before returning from the handler.`
`77`	`78`	`// This keeps the HTTP response open until the final response message is sent.`
`78`		`- await _responseTcs.Task.WaitAsync(postCancellationToken).ConfigureAwait(false);`
	`79`	`+ await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);`
`79`	`80`
`80`	`81`	`return true;`
`81`	`82`	`}`
`@@ -104,22 +105,22 @@ public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken can`
`104`	`105`
`105`	`106`	`var item = new SseItem<JsonRpcMessage?>(message, SseParser.EventTypeDefault);`
`106`	`107`
`107`		`- if (_sseEventStreamWriter is not null)`
	`108`	`+ if (_storeSseWriter is not null)`
`108`	`109`	`{`
`109`		`- item = await _sseEventStreamWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);`
	`110`	`+ item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);`
`110`	`111`	`}`
`111`	`112`
`112`		`- if (!_originalResponseCompleted)`
	`113`	`+ if (!_httpResponseCompleted)`
`113`	`114`	`{`
`114`	`115`	`// Only write the message to the response if the response has not completed.`
`115`	`116`
`116`	`117`	`try`
`117`	`118`	`{`
`118`		`- await _sseResponseWriter.WriteAsync(item, cancellationToken).ConfigureAwait(false);`
	`119`	`+ await _httpSseWriter.WriteAsync(item, cancellationToken).ConfigureAwait(false);`
`119`	`120`	`}`
`120`	`121`	`catch (Exception ex) when (!cancellationToken.IsCancellationRequested)`
`121`	`122`	`{`
`122`		`- _responseTcs.TrySetException(ex);`
	`123`	`+ _httpResponseTcs.TrySetException(ex);`
`123`	`124`	`}`
`124`	`125`	`}`
`125`	`126`	`}`
`@@ -129,8 +130,8 @@ public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken can`
`129`	`130`	`if ((message is JsonRpcResponse or JsonRpcError) && ((JsonRpcMessageWithId)message).Id == _pendingRequest)`
`130`	`131`	`{`
`131`	`132`	`_finalResponseMessageSent = true;`
`132`		`- _responseTcs.TrySetResult(true);`
`133`		`- _streamTcs?.TrySetResult(true);`
	`133`	`+ _httpResponseTcs.TrySetResult(true);`
	`134`	`+ _storeStreamTcs?.TrySetResult(true);`
`134`	`135`	`}`
`135`	`136`	`}`
`136`	`137`	`}`
`@@ -144,61 +145,61 @@ public async ValueTask EnablePollingAsync(TimeSpan retryInterval, CancellationTo`
`144`	`145`
`145`	`146`	`using var _ = await _messageLock.LockAsync(cancellationToken).ConfigureAwait(false);`
`146`	`147`
`147`		`- if (_sseEventStreamWriter is null)`
	`148`	`+ if (_storeSseWriter is null)`
`148`	`149`	`{`
`149`	`150`	`throw new InvalidOperationException($"Polling requires an event stream store to be configured.");`
`150`	`151`	`}`
`151`	`152`
`152`	`153`	`// Send the priming event with the new retry interval.`
`153`		`- var primingItem = await _sseEventStreamWriter.WriteEventAsync(`
	`154`	`+ var primingItem = await _storeSseWriter.WriteEventAsync(`
`154`	`155`	`sseItem: new SseItem<JsonRpcMessage?>() { ReconnectionInterval = retryInterval },`
`155`	`156`	`cancellationToken)`
`156`	`157`	`.ConfigureAwait(false);`
`157`	`158`
`158`	`159`	`// Write to the response stream if it still exists.`
`159`		`- if (!_originalResponseCompleted)`
	`160`	`+ if (!_httpResponseCompleted)`
`160`	`161`	`{`
`161`		`- await _sseResponseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);`
	`162`	`+ await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);`
`162`	`163`	`}`
`163`	`164`
`164`	`165`	`// Set the mode to 'Polling' so that the replay stream ends as soon as all available messages have been sent.`
`165`	`166`	`// This prevents the client from immediately establishing another long-lived connection.`
`166`		`- await _sseEventStreamWriter.SetModeAsync(SseEventStreamMode.Polling, cancellationToken).ConfigureAwait(false);`
	`167`	`+ await _storeSseWriter.SetModeAsync(SseEventStreamMode.Polling, cancellationToken).ConfigureAwait(false);`
`167`	`168`
`168`	`169`	`// Signal completion so HandlePostAsync can return.`
`169`		`- _responseTcs.TrySetResult(true);`
	`170`	`+ _httpResponseTcs.TrySetResult(true);`
`170`	`171`	`}`
`171`	`172`
`172`		`- private async ValueTask<SseItem<JsonRpcMessage?>?> TryStartSseEventStreamAsync(RequestId requestId, CancellationToken cancellationToken)`
	`173`	`+ private async ValueTask<SseItem<JsonRpcMessage?>?> TryStartSseEventStreamAsync(RequestId requestId)`
`173`	`174`	`{`
`174`		`- Debug.Assert(_sseEventStreamWriter is null);`
	`175`	`+ Debug.Assert(_storeSseWriter is null);`
`175`	`176`
`176`		`- _sseEventStreamWriter = await parentTransport.TryCreateEventStreamAsync(`
	`177`	`+ _storeSseWriter = await parentTransport.TryCreateEventStreamAsync(`
`177`	`178`	`streamId: requestId.Id!.ToString()!,`
`178`		`- cancellationToken: cancellationToken)`
	`179`	`+ cancellationToken: sessionCancellationToken)`
`179`	`180`	`.ConfigureAwait(false);`
`180`	`181`
`181`		`- if (_sseEventStreamWriter is null)`
	`182`	`+ if (_storeSseWriter is null)`
`182`	`183`	`{`
`183`	`184`	`return null;`
`184`	`185`	`}`
`185`	`186`
`186`		`- _streamTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);`
`187`		`- _ = HandleStreamWriterDisposalAsync(_streamTcs.Task, cancellationToken);`
	`187`	`+ _storeStreamTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);`
	`188`	`+ _ = HandleStoreStreamDisposalAsync(_storeStreamTcs.Task);`
`188`	`189`
`189`		`- return await _sseEventStreamWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);`
	`190`	`+ return await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), sessionCancellationToken).ConfigureAwait(false);`
`190`	`191`
`191`		`- async Task HandleStreamWriterDisposalAsync(Task streamTask, CancellationToken cancellationToken)`
	`192`	`+ async Task HandleStoreStreamDisposalAsync(Task streamTask)`
`192`	`193`	`{`
`193`	`194`	`try`
`194`	`195`	`{`
`195`		`- await streamTask.WaitAsync(cancellationToken).ConfigureAwait(false);`
	`196`	`+ await streamTask.WaitAsync(sessionCancellationToken).ConfigureAwait(false);`
`196`	`197`	`}`
`197`	`198`	`finally`
`198`	`199`	`{`
`199`	`200`	`using var _ = await _messageLock.LockAsync().ConfigureAwait(false);`
`200`	`201`
`201`		`- await _sseEventStreamWriter!.DisposeAsync().ConfigureAwait(false);`
	`202`	`+ await _storeSseWriter!.DisposeAsync().ConfigureAwait(false);`
`202`	`203`	`}`
`203`	`204`	`}`
`204`	`205`	`}`
`@@ -207,16 +208,16 @@ public async ValueTask DisposeAsync()`
`207`	`208`	`{`
`208`	`209`	`using var _ = await _messageLock.LockAsync().ConfigureAwait(false);`
`209`	`210`
`210`		`- if (_originalResponseCompleted)`
	`211`	`+ if (_httpResponseCompleted)`
`211`	`212`	`{`
`212`	`213`	`return;`
`213`	`214`	`}`
`214`	`215`
`215`		`- _originalResponseCompleted = true;`
	`216`	`+ _httpResponseCompleted = true;`
`216`	`217`
`217`		`- _responseTcs.TrySetResult(true);`
	`218`	`+ _httpResponseTcs.TrySetResult(true);`
`218`	`219`
`219`		`- _sseResponseWriter.Dispose();`
	`220`	`+ _httpSseWriter.Dispose();`
`220`	`221`
`221`	`222`	`// Don't dispose the event stream writer here, as we may continue to write to the event store`
`222`	`223`	`// after disposal if there are pending messages.`
Original file line number	Diff line number	Diff line change
`@@ -36,11 +36,11 @@ public sealed class StreamableHttpServerTransport : ITransport`
`36`	`36`	`private readonly CancellationTokenSource _transportDisposedCts = new();`
`37`	`37`	`private readonly SemaphoreSlim _unsolicitedMessageLock = new(1, 1);`
`38`	`38`
`39`		`- private SseEventWriter? _sseResponseWriter;`
`40`		`- private ISseEventStreamWriter? _sseEventStreamWriter;`
`41`		`- private TaskCompletionSource<bool>? _responseTcs;`
`42`		`- private bool _getRequestStarted;`
`43`		`- private bool _originalResponseCompleted;`
	`39`	`+ private SseEventWriter? _httpSseWriter;`
	`40`	`+ private ISseEventStreamWriter? _storeSseWriter;`
	`41`	`+ private TaskCompletionSource<bool>? _httpResponseTcs;`
	`42`	`+ private bool _getHttpRequestStarted;`
	`43`	`+ private bool _getHttpResponseCompleted;`
`44`	`44`
`45`	`45`	`/// <inheritdoc/>`
`46`	`46`	`public string? SessionId { get; init; }`
`@@ -111,25 +111,25 @@ public async Task HandleGetRequestAsync(Stream sseResponseStream, CancellationTo`
`111`	`111`
`112`	`112`	`using (await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false))`
`113`	`113`	`{`
`114`		`- if (_getRequestStarted)`
	`114`	`+ if (_getHttpRequestStarted)`
`115`	`115`	`{`
`116`	`116`	`throw new InvalidOperationException("Session resumption is not yet supported. Please start a new session.");`
`117`	`117`	`}`
`118`	`118`
`119`		`- _getRequestStarted = true;`
`120`		`- _sseResponseWriter = new SseEventWriter(sseResponseStream);`
`121`		`- _responseTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);`
`122`		`- _sseEventStreamWriter = await TryCreateEventStreamAsync(streamId: UnsolicitedMessageStreamId, cancellationToken).ConfigureAwait(false);`
`123`		`- if (_sseEventStreamWriter is not null)`
	`119`	`+ _getHttpRequestStarted = true;`
	`120`	`+ _httpSseWriter = new SseEventWriter(sseResponseStream);`
	`121`	`+ _httpResponseTcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);`
	`122`	`+ _storeSseWriter = await TryCreateEventStreamAsync(streamId: UnsolicitedMessageStreamId, cancellationToken).ConfigureAwait(false);`
	`123`	`+ if (_storeSseWriter is not null)`
`124`	`124`	`{`
`125`		`- var primingItem = await _sseEventStreamWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);`
`126`		`- await _sseResponseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);`
	`125`	`+ var primingItem = await _storeSseWriter.WriteEventAsync(SseItem.Prime<JsonRpcMessage>(), cancellationToken).ConfigureAwait(false);`
	`126`	`+ await _httpSseWriter.WriteAsync(primingItem, cancellationToken).ConfigureAwait(false);`
`127`	`127`	`}`
`128`	`128`	`}`
`129`	`129`
`130`	`130`	`// Wait for the response to be written before returning from the handler.`
`131`	`131`	`// This keeps the HTTP response open until the final response message is sent.`
`132`		`- await _responseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);`
	`132`	`+ await _httpResponseTcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);`
`133`	`133`	`}`
`134`	`134`
`135`	`135`	`/// <summary>`
`@@ -155,14 +155,13 @@ public async Task<bool> HandlePostRequestAsync(JsonRpcMessage message, Stream re`
`155`	`155`	`Throw.IfNull(message);`
`156`	`156`	`Throw.IfNull(responseStream);`
`157`	`157`
	`158`	`+ var postTransport = new StreamableHttpPostTransport(this, responseStream, _transportDisposedCts.Token);`
`158`	`159`	`using var postCts = CancellationTokenSource.CreateLinkedTokenSource(_transportDisposedCts.Token, cancellationToken);`
`159`		`- var postTransport = new StreamableHttpPostTransport(this, responseStream);`
`160`	`160`	`await using (postTransport.ConfigureAwait(false))`
`161`	`161`	`{`
`162`	`162`	`return await postTransport.HandlePostAsync(`
`163`	`163`	`message,`
`164`		`- postCancellationToken: postCts.Token,`
`165`		`- sessionCancellationToken: _transportDisposedCts.Token)`
	`164`	`+ cancellationToken: postCts.Token)`
`166`	`165`	`.ConfigureAwait(false);`
`167`	`166`	`}`
`168`	`167`	`}`
`@@ -179,34 +178,34 @@ public async Task SendMessageAsync(JsonRpcMessage message, CancellationToken can`
`179`	`178`
`180`	`179`	`using var _ = await _unsolicitedMessageLock.LockAsync(cancellationToken).ConfigureAwait(false);`
`181`	`180`
`182`		`- if (!_getRequestStarted)`
	`181`	`+ if (!_getHttpRequestStarted)`
`183`	`182`	`{`
`184`	`183`	`// Clients are not required to make a GET request for unsolicited messages.`
`185`	`184`	`// If no GET request has been made, drop the message.`
`186`	`185`	`return;`
`187`	`186`	`}`
`188`	`187`
`189`		`- Debug.Assert(_sseResponseWriter is not null);`
`190`		`- Debug.Assert(_responseTcs is not null);`
	`188`	`+ Debug.Assert(_httpSseWriter is not null);`
	`189`	`+ Debug.Assert(_httpResponseTcs is not null);`
`191`	`190`
`192`	`191`	`var item = SseItem.Message(message);`
`193`	`192`
`194`		`- if (_sseEventStreamWriter is not null)`
	`193`	`+ if (_storeSseWriter is not null)`
`195`	`194`	`{`
`196`		`- item = await _sseEventStreamWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);`
	`195`	`+ item = await _storeSseWriter.WriteEventAsync(item, cancellationToken).ConfigureAwait(false);`
`197`	`196`	`}`
`198`	`197`
`199`		`- if (!_originalResponseCompleted)`
	`198`	`+ if (!_getHttpResponseCompleted)`
`200`	`199`	`{`
`201`	`200`	`// Only write the message to the response if the response has not completed.`
`202`	`201`
`203`	`202`	`try`
`204`	`203`	`{`
`205`		`- await _sseResponseWriter!.WriteAsync(item, cancellationToken).ConfigureAwait(false);`
	`204`	`+ await _httpSseWriter!.WriteAsync(item, cancellationToken).ConfigureAwait(false);`
`206`	`205`	`}`
`207`	`206`	`catch (Exception ex) when (!cancellationToken.IsCancellationRequested)`
`208`	`207`	`{`
`209`		`- _responseTcs!.TrySetException(ex);`
	`208`	`+ _httpResponseTcs!.TrySetException(ex);`
`210`	`209`	`}`
`211`	`210`	`}`
`212`	`211`	`}`
`@@ -216,12 +215,12 @@ public async ValueTask DisposeAsync()`
`216`	`215`	`{`
`217`	`216`	`using var _ = await _unsolicitedMessageLock.LockAsync().ConfigureAwait(false);`
`218`	`217`
`219`		`- if (_originalResponseCompleted)`
	`218`	`+ if (_getHttpResponseCompleted)`
`220`	`219`	`{`
`221`	`220`	`return;`
`222`	`221`	`}`
`223`	`222`
`224`		`- _originalResponseCompleted = true;`
	`223`	`+ _getHttpResponseCompleted = true;`
`225`	`224`
`226`	`225`	`try`
`227`	`226`	`{`
`@@ -232,12 +231,12 @@ public async ValueTask DisposeAsync()`
`232`	`231`	`{`
`233`	`232`	`try`
`234`	`233`	`{`
`235`		`- _responseTcs?.TrySetResult(true);`
`236`		`- _sseResponseWriter?.Dispose();`
	`234`	`+ _httpResponseTcs?.TrySetResult(true);`
	`235`	`+ _httpSseWriter?.Dispose();`
`237`	`236`
`238`		`- if (_sseEventStreamWriter is not null)`
	`237`	`+ if (_storeSseWriter is not null)`
`239`	`238`	`{`
`240`		`- await _sseEventStreamWriter.DisposeAsync().ConfigureAwait(false);`
	`239`	`+ await _storeSseWriter.DisposeAsync().ConfigureAwait(false);`
`241`	`240`	`}`
`242`	`241`	`}`
`243`	`242`	`finally`