Skip to content

Commit 860e132

Browse files
committed
Add HttpClientCompletionDetails for HTTP session completion
1 parent 0b8138c commit 860e132

File tree

9 files changed

+267
-20
lines changed

9 files changed

+267
-20
lines changed
Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
using System.Net;
2+
3+
namespace ModelContextProtocol.Client;
4+
5+
/// <summary>
6+
/// Provides details about the completion of an HTTP-based MCP client session,
7+
/// including sessions using the legacy SSE transport or the Streamable HTTP transport.
8+
/// </summary>
9+
public sealed class HttpClientCompletionDetails : ClientCompletionDetails
10+
{
11+
/// <summary>
12+
/// Gets the HTTP status code that caused the session to close, or <see langword="null"/> if unavailable.
13+
/// </summary>
14+
public HttpStatusCode? HttpStatusCode { get; set; }
15+
}

src/ModelContextProtocol.Core/Client/SseClientSessionTransport.cs

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
using Microsoft.Extensions.Logging.Abstractions;
33
using ModelContextProtocol.Protocol;
44
using System.Diagnostics;
5+
using System.Net;
56
using System.Net.Http.Headers;
67
using System.Net.ServerSentEvents;
78
using System.Text.Json;
@@ -124,7 +125,7 @@ private async Task CloseAsync()
124125
}
125126
finally
126127
{
127-
SetDisconnected();
128+
SetDisconnected(new TransportClosedException(new HttpClientCompletionDetails()));
128129
}
129130
}
130131

@@ -143,6 +144,7 @@ public override async ValueTask DisposeAsync()
143144

144145
private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
145146
{
147+
HttpStatusCode? failureStatusCode = null;
146148
try
147149
{
148150
using var request = new HttpRequestMessage(HttpMethod.Get, _sseEndpoint);
@@ -151,6 +153,11 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
151153

152154
using var response = await _httpClient.SendAsync(request, message: null, cancellationToken).ConfigureAwait(false);
153155

156+
if (!response.IsSuccessStatusCode)
157+
{
158+
failureStatusCode = response.StatusCode;
159+
}
160+
154161
await response.EnsureSuccessStatusCodeWithResponseBodyAsync(cancellationToken).ConfigureAwait(false);
155162

156163
using var stream = await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false);
@@ -179,7 +186,11 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
179186
}
180187
else
181188
{
182-
SetDisconnected(ex);
189+
SetDisconnected(new TransportClosedException(new HttpClientCompletionDetails
190+
{
191+
HttpStatusCode = failureStatusCode,
192+
Exception = ex,
193+
}));
183194

184195
LogTransportReadMessagesFailed(Name, ex);
185196
_connectionEstablished.TrySetException(ex);
@@ -188,7 +199,7 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
188199
}
189200
finally
190201
{
191-
SetDisconnected();
202+
SetDisconnected(new TransportClosedException(new HttpClientCompletionDetails()));
192203
}
193204
}
194205

src/ModelContextProtocol.Core/Client/StreamableHttpClientSessionTransport.cs

Lines changed: 38 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,12 @@ internal sealed partial class StreamableHttpClientSessionTransport : TransportBa
2020

2121
private readonly McpHttpClient _httpClient;
2222
private readonly HttpClientTransportOptions _options;
23-
private readonly CancellationTokenSource _connectionCts;
23+
private readonly CancellationTokenSource _connectionCts = new();
2424
private readonly ILogger _logger;
2525

2626
private string? _negotiatedProtocolVersion;
2727
private Task? _getReceiveTask;
28+
private volatile TransportClosedException? _disconnectError;
2829

2930
private readonly SemaphoreSlim _disposeLock = new(1, 1);
3031
private bool _disposed;
@@ -42,7 +43,6 @@ public StreamableHttpClientSessionTransport(
4243

4344
_options = transportOptions;
4445
_httpClient = httpClient;
45-
_connectionCts = new CancellationTokenSource();
4646
_logger = (ILogger?)loggerFactory?.CreateLogger<HttpClientTransport>() ?? NullLogger.Instance;
4747

4848
// We connect with the initialization request with the MCP transport. This means that any errors won't be observed
@@ -96,6 +96,13 @@ internal async Task<HttpResponseMessage> SendHttpRequestAsync(JsonRpcMessage mes
9696
// We'll let the caller decide whether to throw or fall back given an unsuccessful response.
9797
if (!response.IsSuccessStatusCode)
9898
{
99+
// Per the MCP spec, a 404 response to a request containing an Mcp-Session-Id
100+
// indicates the session has ended. Signal completion so McpClient.Completion resolves.
101+
if (response.StatusCode == HttpStatusCode.NotFound && SessionId is not null)
102+
{
103+
SetSessionExpired(response.StatusCode);
104+
}
105+
99106
return response;
100107
}
101108

@@ -184,18 +191,14 @@ public override async ValueTask DisposeAsync()
184191
{
185192
LogTransportShutdownFailed(Name, ex);
186193
}
187-
finally
188-
{
189-
_connectionCts.Dispose();
190-
}
191194
}
192195
finally
193196
{
194197
// If we're auto-detecting the transport and failed to connect, leave the message Channel open for the SSE transport.
195198
// This class isn't directly exposed to public callers, so we don't have to worry about changing the _state in this case.
196199
if (_options.TransportMode is not HttpTransportMode.AutoDetect || _getReceiveTask is not null)
197200
{
198-
SetDisconnected();
201+
SetDisconnected(_disconnectError ?? new TransportClosedException(new HttpClientCompletionDetails()));
199202
}
200203
}
201204
}
@@ -204,8 +207,8 @@ private async Task ReceiveUnsolicitedMessagesAsync()
204207
{
205208
var state = new SseStreamState();
206209

207-
// Continuously receive unsolicited messages until canceled
208-
while (!_connectionCts.Token.IsCancellationRequested)
210+
// Continuously receive unsolicited messages until canceled or disconnected
211+
while (!_connectionCts.Token.IsCancellationRequested && IsConnected)
209212
{
210213
await SendGetSseRequestWithRetriesAsync(
211214
relatedRpcRequest: null,
@@ -285,6 +288,13 @@ await SendGetSseRequestWithRetriesAsync(
285288

286289
if (!response.IsSuccessStatusCode)
287290
{
291+
// Per the MCP spec, a 404 response to a request containing an Mcp-Session-Id
292+
// indicates the session has ended. Signal completion so McpClient.Completion resolves.
293+
if (response.StatusCode == HttpStatusCode.NotFound && SessionId is not null)
294+
{
295+
SetSessionExpired(response.StatusCode);
296+
}
297+
288298
// If the server could be reached but returned a non-success status code,
289299
// retrying likely won't change that.
290300
return null;
@@ -474,4 +484,23 @@ private static TimeSpan ElapsedSince(long stopwatchTimestamp)
474484
return TimeSpan.FromSeconds((double)(Stopwatch.GetTimestamp() - stopwatchTimestamp) / Stopwatch.Frequency);
475485
#endif
476486
}
487+
488+
private void SetSessionExpired(HttpStatusCode statusCode)
489+
{
490+
// Store the error before canceling so DisposeAsync can use it if it races us, especially
491+
// after the call to Cancel below, to invoke SetDisconnected.
492+
_disconnectError = new TransportClosedException(new HttpClientCompletionDetails
493+
{
494+
HttpStatusCode = statusCode,
495+
Exception = new McpException(
496+
"The server returned HTTP 404 for a request with an Mcp-Session-Id, indicating the session has expired. " +
497+
"To continue, create a new client session or call ResumeSessionAsync with a new connection."),
498+
});
499+
500+
// Cancel to unblock any in-flight operations (e.g., SSE stream reads in
501+
// SendGetSseRequestWithRetriesAsync) that are waiting on _connectionCts.Token.
502+
_connectionCts.Cancel();
503+
504+
SetDisconnected(_disconnectError);
505+
}
477506
}

src/ModelContextProtocol.Core/McpSessionHandler.cs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -909,9 +909,11 @@ public async ValueTask DisposeAsync()
909909
{
910910
await _messageProcessingTask.ConfigureAwait(false);
911911
}
912-
catch (OperationCanceledException)
912+
catch
913913
{
914-
// Ignore cancellation
914+
// Ignore exceptions from the message processing loop. It may fault with
915+
// OperationCanceledException on normal shutdown or TransportClosedException
916+
// when the transport's channel completes with an error.
915917
}
916918
}
917919

src/ModelContextProtocol.Core/Protocol/TransportBase.cs

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,13 +90,15 @@ internal TransportBase(string name, Channel<JsonRpcMessage>? messageChannel, ILo
9090
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
9191
protected async Task WriteMessageAsync(JsonRpcMessage message, CancellationToken cancellationToken = default)
9292
{
93+
cancellationToken.ThrowIfCancellationRequested();
94+
9395
if (!IsConnected)
9496
{
95-
throw new InvalidOperationException("Transport is not connected.");
97+
// Transport disconnected concurrently. Silently drop rather than throw,
98+
// to avoid surfacing spurious errors during shutdown races.
99+
return;
96100
}
97101

98-
cancellationToken.ThrowIfCancellationRequested();
99-
100102
if (_logger.IsEnabled(LogLevel.Debug))
101103
{
102104
var messageId = (message as JsonRpcMessageWithId)?.Id.ToString() ?? "(no id)";

tests/ModelContextProtocol.AspNetCore.Tests/HttpServerIntegrationTests.cs

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,8 @@ public async Task Completion_GracefulDisposal_ReturnsCompletionDetails()
314314
Assert.True(client.Completion.IsCompleted);
315315

316316
var details = await client.Completion;
317-
Assert.Null(details.Exception);
317+
var httpDetails = Assert.IsType<HttpClientCompletionDetails>(details);
318+
Assert.Null(httpDetails.Exception);
319+
Assert.Null(httpDetails.HttpStatusCode);
318320
}
319321
}

tests/ModelContextProtocol.AspNetCore.Tests/SseIntegrationTests.cs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -322,7 +322,8 @@ public async Task Completion_ServerShutdown_ReturnsHttpCompletionDetails()
322322
await app.StopAsync(TestContext.Current.CancellationToken);
323323

324324
var details = await mcpClient.Completion.WaitAsync(TestContext.Current.CancellationToken);
325-
Assert.IsType<ClientCompletionDetails>(details);
325+
var httpDetails = Assert.IsType<HttpClientCompletionDetails>(details);
326+
Assert.Null(httpDetails.HttpStatusCode);
326327
}
327328

328329
public class Envelope

0 commit comments

Comments
 (0)