@@ -694,23 +694,14 @@ public async Task NotifyCancel_Should_Be_Handled()
694694 {
695695 // Arrange
696696 TaskCompletionSource < JsonRpcNotification > notificationReceived = new ( ) ;
697- TaskCompletionSource notificationIntercepted = new ( ) ;
698697 await using TestServerTransport transport = new ( LoggerFactory ) ;
699698 transport . OnMessageSent = ( message ) =>
700699 {
701700 if ( message is JsonRpcNotification notification
702701 && notification . Method == NotificationMethods . CancelledNotification )
703702 notificationReceived . TrySetResult ( notification ) ;
704703 } ;
705- var options = CreateOptions ( new ( )
706- {
707- NotificationHandlers = [ new ( NotificationMethods . CancelledNotification , ( notification , cancellationToken ) =>
708- {
709- InvalidOperationException exception = new ( "The sender of a notification shouldn't handle the notification." ) ;
710- notificationIntercepted . TrySetException ( exception ) ;
711- return notificationIntercepted . Task ;
712- } ) ] ,
713- } ) ;
704+ var options = CreateOptions ( ) ;
714705 await using var server = McpServerFactory . Create ( transport , options , LoggerFactory ) ;
715706
716707 // Act
@@ -732,10 +723,46 @@ await server.SendNotificationAsync(
732723 Assert . NotNull ( cancelled ) ;
733724 Assert . Equal ( "abc" , cancelled . RequestId . ToString ( ) ) ;
734725 Assert . Equal ( "Cancelled" , cancelled . Reason ) ;
726+ }
735727
736- Assert . Throws < TimeoutException > ( ( ) => notificationIntercepted . Task
737- . Wait ( TimeSpan . FromSeconds ( 5 ) , token ) ) ;
738- Assert . False ( notificationIntercepted . Task . IsCompleted ,
739- "Notifications should not be intercepted by the sender of the notification." ) ;
728+ [ Fact ]
729+ public async Task SendRequest_Should_Notify_When_Cancelled ( )
730+ {
731+ // Arrange
732+ TaskCompletionSource < JsonRpcNotification > notificationReceived = new ( ) ;
733+ await using TestServerTransport transport = new ( LoggerFactory ) ;
734+ transport . OnMessageSent = ( message ) =>
735+ {
736+ if ( message is JsonRpcNotification notification
737+ && notification . Method == NotificationMethods . CancelledNotification )
738+ notificationReceived . TrySetResult ( notification ) ;
739+ } ;
740+ var options = CreateOptions ( ) ;
741+ await using var server = McpServerFactory . Create ( transport , options , LoggerFactory ) ;
742+
743+ // Act
744+ var token = TestContext . Current . CancellationToken ;
745+ Task serverTask = server . RunAsync ( token ) ;
746+ using CancellationTokenSource cts = new ( ) ;
747+ await cts . CancelAsync ( ) ;
748+
749+ await Assert . ThrowsAsync < TaskCanceledException > ( async ( ) =>
750+ {
751+ await server . SendRequestAsync ( new JsonRpcRequest
752+ {
753+ Method = RequestMethods . Ping ,
754+ Id = new ( "abc" ) ,
755+ } , cts . Token ) ;
756+ } ) ;
757+
758+ await server . DisposeAsync ( ) ;
759+ await serverTask . WaitAsync ( TimeSpan . FromSeconds ( 1 ) , token ) ;
760+ var notification = await notificationReceived . Task . WaitAsync ( TimeSpan . FromSeconds ( 1 ) , token ) ;
761+
762+ // Assert
763+ var cancelled = JsonSerializer . Deserialize < CancelledNotification > ( notification . Params ) ;
764+ Assert . NotNull ( cancelled ) ;
765+ Assert . Equal ( "abc" , cancelled . RequestId . ToString ( ) ) ;
766+ Assert . Null ( cancelled . Reason ) ;
740767 }
741768}
0 commit comments