Skip to content

Commit 947441b

Browse files
Copilotstephentoub
andcommitted
Prevent cancellation of initialize request per MCP spec
Per MCP spec: "The initialize request MUST NOT be cancelled by clients" Client side: Skip RegisterCancellation in SendRequestAsync for initialize requests so that no notifications/cancelled is sent even if the cancellation token fires (e.g. from timeout). Server side: Don't store initialize request in _handlingRequests so incoming cancellation notifications can't cancel it. Added tests for both behaviors. Co-authored-by: stephentoub <2642209+stephentoub@users.noreply.github.com>
1 parent f0a21f7 commit 947441b

3 files changed

Lines changed: 99 additions & 3 deletions

File tree

src/ModelContextProtocol.Core/McpSessionHandler.cs

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -190,11 +190,16 @@ async Task ProcessMessageAsync()
190190
{
191191
// Register before we yield, so that the tracking is guaranteed to be there
192192
// when subsequent messages arrive, even if the asynchronous processing happens
193-
// out of order.
193+
// out of order. Per spec, "The initialize request MUST NOT be cancelled by clients",
194+
// so we don't track it in _handlingRequests to prevent cancellation notifications from
195+
// canceling it.
194196
if (messageWithId is not null)
195197
{
196198
combinedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
197-
_handlingRequests[messageWithId.Id] = combinedCts;
199+
if (message is not JsonRpcRequest { Method: RequestMethods.Initialize })
200+
{
201+
_handlingRequests[messageWithId.Id] = combinedCts;
202+
}
198203
}
199204

200205
// If we await the handler without yielding first, the transport may not be able to read more messages,
@@ -528,9 +533,10 @@ public async Task<JsonRpcResponse> SendRequestAsync(JsonRpcRequest request, Canc
528533
// Now that the request has been sent, register for cancellation. If we registered before,
529534
// a cancellation request could arrive before the server knew about that request ID, in which
530535
// case the server could ignore it.
536+
// Per spec, "The initialize request MUST NOT be cancelled by clients", so skip registration for initialize.
531537
LogRequestSentAwaitingResponse(EndpointName, request.Method, request.Id);
532538
JsonRpcMessage? response;
533-
using (var registration = RegisterCancellation(cancellationToken, request))
539+
using (var registration = method != RequestMethods.Initialize ? RegisterCancellation(cancellationToken, request) : default)
534540
{
535541
response = await tcs.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
536542
}

tests/ModelContextProtocol.Tests/Protocol/CancellationTests.cs

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
using Microsoft.Extensions.DependencyInjection;
2+
using ModelContextProtocol.Client;
23
using ModelContextProtocol.Protocol;
34
using ModelContextProtocol.Server;
5+
using ModelContextProtocol.Tests.Utils;
6+
using System.IO.Pipelines;
7+
using System.Text;
48

59
namespace ModelContextProtocol.Tests;
610

@@ -65,4 +69,37 @@ public async Task CancellationPropagation_RequestingCancellationCancelsPendingRe
6569
cts.Cancel();
6670
await Assert.ThrowsAnyAsync<OperationCanceledException>(async () => await waitTask);
6771
}
72+
73+
[Fact]
74+
public async Task InitializeTimeout_DoesNotSendCancellationNotification()
75+
{
76+
// Arrange: Create a transport where the server never responds, so the client will time out.
77+
var serverInput = new MemoryStream();
78+
var serverOutputPipe = new Pipe();
79+
80+
var clientTransport = new StreamClientTransport(
81+
serverInput: serverInput,
82+
serverOutputPipe.Reader.AsStream(),
83+
LoggerFactory);
84+
85+
var clientOptions = new McpClientOptions
86+
{
87+
InitializationTimeout = TimeSpan.FromMilliseconds(500),
88+
};
89+
90+
// Act: Client will send initialize, then time out since no response comes.
91+
// Per spec, "The initialize request MUST NOT be cancelled by clients",
92+
// so no cancellation notification should be sent.
93+
await Assert.ThrowsAsync<TimeoutException>(async () =>
94+
{
95+
await McpClient.CreateAsync(clientTransport, clientOptions: clientOptions, loggerFactory: LoggerFactory,
96+
cancellationToken: TestContext.Current.CancellationToken);
97+
});
98+
99+
// Assert: Read what was written to serverInput.
100+
// The only message should be the initialize request, NOT a cancellation notification.
101+
var content = Encoding.UTF8.GetString(serverInput.ToArray());
102+
Assert.Contains("\"method\":\"initialize\"", content);
103+
Assert.DoesNotContain("notifications/cancelled", content);
104+
}
68105
}

tests/ModelContextProtocol.Tests/Server/McpServerTests.cs

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -861,6 +861,59 @@ public async Task Can_SendMessage_Before_RunAsync()
861861
Assert.Same(logNotification, transport.SentMessages[0]);
862862
}
863863

864+
[Fact]
865+
public async Task Server_IgnoresCancellationNotificationForInitializeRequest()
866+
{
867+
// Arrange
868+
await using var transport = new TestServerTransport();
869+
await using McpServer server = McpServer.Create(transport, _options, LoggerFactory);
870+
var runTask = server.RunAsync(TestContext.Current.CancellationToken);
871+
872+
// Set up to capture the initialize response
873+
var initializeRequest = new JsonRpcRequest
874+
{
875+
Id = new RequestId("init-cancel-test"),
876+
Method = RequestMethods.Initialize,
877+
Params = JsonSerializer.SerializeToNode(new InitializeRequestParams
878+
{
879+
ProtocolVersion = "2024-11-05",
880+
Capabilities = new ClientCapabilities(),
881+
ClientInfo = new Implementation { Name = "test-client", Version = "1.0.0" }
882+
}, McpJsonUtilities.DefaultOptions)
883+
};
884+
885+
var initResponseTcs = new TaskCompletionSource<JsonRpcResponse>();
886+
transport.OnMessageSent = (message) =>
887+
{
888+
if (message is JsonRpcResponse response && response.Id == initializeRequest.Id)
889+
{
890+
initResponseTcs.TrySetResult(response);
891+
}
892+
};
893+
894+
// Act: Send initialize request and immediately send a cancellation notification for it.
895+
// Per spec, "The initialize request MUST NOT be cancelled by clients", so the server
896+
// should ignore the cancellation and still complete the initialize request.
897+
await transport.SendClientMessageAsync(initializeRequest, TestContext.Current.CancellationToken);
898+
await transport.SendClientMessageAsync(new JsonRpcNotification
899+
{
900+
Method = NotificationMethods.CancelledNotification,
901+
Params = JsonSerializer.SerializeToNode(
902+
new CancelledNotificationParams { RequestId = initializeRequest.Id },
903+
McpJsonUtilities.DefaultOptions),
904+
}, TestContext.Current.CancellationToken);
905+
906+
// Assert: The initialize response should still arrive (not cancelled)
907+
var response = await initResponseTcs.Task.WaitAsync(TestConstants.DefaultTimeout, TestContext.Current.CancellationToken);
908+
Assert.NotNull(response.Result);
909+
var initResult = JsonSerializer.Deserialize<InitializeResult>(response.Result, McpJsonUtilities.DefaultOptions);
910+
Assert.NotNull(initResult);
911+
Assert.NotNull(initResult.ServerInfo);
912+
913+
await transport.DisposeAsync();
914+
await runTask;
915+
}
916+
864917
private static async Task InitializeServerAsync(TestServerTransport transport, ClientCapabilities capabilities, CancellationToken cancellationToken = default)
865918
{
866919
var initializeRequest = new JsonRpcRequest

0 commit comments

Comments
 (0)