1919import java .util .Queue ;
2020import java .util .concurrent .CompletableFuture ;
2121import java .util .concurrent .ConcurrentLinkedQueue ;
22+ import java .util .concurrent .TimeUnit ;
2223import java .util .concurrent .atomic .AtomicBoolean ;
2324import java .util .concurrent .atomic .AtomicReference ;
2425import java .util .function .Function ;
3940import io .modelcontextprotocol .spec .HttpHeaders ;
4041import io .modelcontextprotocol .spec .McpSchema ;
4142import io .modelcontextprotocol .util .KeepAliveScheduler ;
43+ import io .modelcontextprotocol .util .ToolsUtils ;
4244import jakarta .servlet .http .HttpServletRequest ;
4345import jakarta .servlet .http .HttpServletResponse ;
4446import org .apache .catalina .LifecycleException ;
5557import org .slf4j .LoggerFactory ;
5658import reactor .core .publisher .Flux ;
5759import reactor .core .publisher .Mono ;
60+ import reactor .core .publisher .Sinks ;
5861import reactor .core .scheduler .Schedulers ;
5962import reactor .test .StepVerifier ;
6063import reactor .test .scheduler .VirtualTimeScheduler ;
@@ -205,8 +208,9 @@ void testMissingHandlerReturnsMethodNotFoundError() {
205208 }
206209
207210 @ Test
208- void listeningStreamHasStreamingHeaders () throws Exception {
209- prepareAsyncServerBuilder ().serverInfo ("test-server" , "1.0.0" ).build ();
211+ void listeningStreamSendsHeadersBeforeFirstEvent () throws Exception {
212+ // No keep-alive, so the server never writes to the stream
213+ withTransportProvider (null , null );
210214 var sessionId = initializeSession (httpClient );
211215
212216 var get = HttpRequest .newBuilder ()
@@ -216,8 +220,10 @@ void listeningStreamHasStreamingHeaders() throws Exception {
216220 .GET ()
217221 .build ();
218222
219- var response = httpClient .send (get , HttpResponse .BodyHandlers .ofInputStream ());
223+ // The response completes as soon as the headers are received
224+ var response = httpClient .sendAsync (get , HttpResponse .BodyHandlers .ofInputStream ()).get (5 , TimeUnit .SECONDS );
220225 try {
226+ assertThat (response .statusCode ()).isEqualTo (HttpServletResponse .SC_OK );
221227 assertSseHeaders (response );
222228 }
223229 finally {
@@ -226,17 +232,29 @@ void listeningStreamHasStreamingHeaders() throws Exception {
226232 }
227233
228234 @ Test
229- void toolCallSseResponseHasStreamingHeaders () throws Exception {
230- prepareAsyncServerBuilder ().serverInfo ("test-server" , "1.0.0" )
231- .capabilities (McpSchema .ServerCapabilities .builder ().tools (true ).build ())
235+ void toolCallSseResponseSendsHeadersBeforeResult () throws Exception {
236+ mcpServerTransportProvider .closeGracefully ().block ();
237+ mcpServerTransportProvider = HttpServletStreamableServerTransportProvider .builder ()
238+ .contextExtractor (TEST_CONTEXT_EXTRACTOR )
239+ .mcpEndpoint (MESSAGE_ENDPOINT )
240+ // No keep-alive, so the server writes nothing until the tool returns
241+ .keepAliveInterval (null )
242+ // no sweeping: the session must survive every keep-alive round
243+ .sessionSweepInterval (null )
244+ .build ();
245+
246+ MCP_SERVLET .setDelegate (mcpServerTransportProvider );
247+
248+ // Custom server with a tool that hangs
249+ Sinks .One <McpSchema .CallToolResult > toolResult = Sinks .one ();
250+ prepareAsyncServerBuilder ().capabilities (McpSchema .ServerCapabilities .builder ().tools (true ).build ())
232251 .tools (McpServerFeatures .AsyncToolSpecification .builder ()
233- .tool (McpSchema .Tool .builder ("echo" , EMPTY_JSON_SCHEMA ).description ("returns immediately" ).build ())
234- .callHandler ((exchange ,
235- request ) -> Mono .just (McpSchema .CallToolResult .builder ()
236- .content (List .of (McpSchema .TextContent .builder ("ok" ).build ()))
237- .isError (false )
238- .build ()))
252+ .tool (McpSchema .Tool .builder ("slow" , ToolsUtils .EMPTY_JSON_SCHEMA )
253+ .description ("returns when released" )
254+ .build ())
255+ .callHandler ((exchange , request ) -> toolResult .asMono ())
239256 .build ())
257+ .serverInfo ("test-server" , "1.0.0" )
240258 .build ();
241259 var sessionId = initializeSession (httpClient );
242260
@@ -247,15 +265,27 @@ void toolCallSseResponseHasStreamingHeaders() throws Exception {
247265 .header (HttpHeaders .MCP_SESSION_ID , sessionId )
248266 .POST (HttpRequest .BodyPublishers
249267 .ofString ("{\" jsonrpc\" :\" 2.0\" ,\" id\" :\" call-1\" ,\" method\" :\" tools/call\" ,"
250- + "\" params\" :{\" name\" :\" echo \" ,\" arguments\" :{}}}" ))
268+ + "\" params\" :{\" name\" :\" slow \" ,\" arguments\" :{}}}" ))
251269 .build ();
252270
253- var response = httpClient .send (post , HttpResponse .BodyHandlers .ofInputStream ());
254271 try {
255- assertSseHeaders (response );
272+ // The response completes as soon as the headers are received
273+ var response = httpClient .sendAsync (post , HttpResponse .BodyHandlers .ofInputStream ())
274+ .get (5 , TimeUnit .SECONDS );
275+ try {
276+ assertThat (response .statusCode ()).isEqualTo (HttpServletResponse .SC_OK );
277+ assertSseHeaders (response );
278+ }
279+ finally {
280+ response .body ().close ();
281+ }
256282 }
257283 finally {
258- response .body ().close ();
284+ // Release the Tomcat thread blocked in doPost until the tool returns
285+ toolResult .tryEmitValue (McpSchema .CallToolResult .builder ()
286+ .content (List .of (McpSchema .TextContent .builder ("done" ).build ()))
287+ .isError (false )
288+ .build ());
259289 }
260290 }
261291
@@ -563,10 +593,6 @@ void sessionEvictionAfterReleasingStream() {
563593 var sessionId = initializeSession (httpClient );
564594 var clientStream = openListeningStream (httpClient , sessionId , null );
565595 awaitClientStreamEstablished ();
566- // "prep" the client so we can close the stream by sending data: the
567- // subscription
568- // in the clientStream is only present when the client has received data
569- awaitNotification (clientStream .events ());
570596 virtualTime .advanceTimeBy (SESSION_SWEEP_INTERVAL .multipliedBy (2 ));
571597 assertThat (postNotification (httpClient , sessionId )).as ("Session should be active when a stream is open" )
572598 .isEqualTo (HttpServletResponse .SC_ACCEPTED );
@@ -621,17 +647,22 @@ void sessionIsEvictedWhenTheClientAbortsAResponseStream() {
621647 try {
622648 withTransportProvider (null , SESSION_SWEEP_INTERVAL );
623649 var toolEmissionInterval = Duration .ofMillis (10 );
650+ var toolCalled = new AtomicBoolean ();
624651 prepareAsyncServerBuilder ().serverInfo ("test-server" , "1.0.0" )
625652 .capabilities (McpSchema .ServerCapabilities .builder ().tools (true ).build ())
626653 .tools (McpServerFeatures .AsyncToolSpecification .builder ()
627654 .tool (McpSchema .Tool .builder ("hangs" , EMPTY_JSON_SCHEMA ).description ("never returns" ).build ())
628- .callHandler ((exchange , request ) -> Flux
629- .interval (toolEmissionInterval , toolEmissionInterval , Schedulers .newSingle ("test-tool" ))
630- .flatMap (tick -> exchange .loggingNotification (McpSchema .LoggingMessageNotification .builder ()
631- .level (McpSchema .LoggingLevel .INFO )
632- .data ("x" .repeat (64 * 1024 ))
633- .build ()))
634- .then (Mono .<McpSchema .CallToolResult >never ()))
655+ .callHandler ((exchange , request ) -> {
656+ toolCalled .set (true );
657+ // Write until the server notices the client is gone
658+ return Flux
659+ .interval (toolEmissionInterval , toolEmissionInterval , Schedulers .newSingle ("test-tool" ))
660+ .flatMap (tick -> exchange .loggingNotification (McpSchema .LoggingMessageNotification .builder ()
661+ .level (McpSchema .LoggingLevel .INFO )
662+ .data ("x" .repeat (64 * 1024 ))
663+ .build ()))
664+ .then (Mono .<McpSchema .CallToolResult >never ());
665+ })
635666 .build ())
636667 .build ();
637668 var sessionId = initializeSession (httpClient );
@@ -640,18 +671,7 @@ void sessionIsEvictedWhenTheClientAbortsAResponseStream() {
640671 // stream for as long as the call is in flight
641672 var responseStream = postToolCall (httpClient , sessionId , "hangs" );
642673
643- await ().atMost (Duration .ofSeconds (5 ))
644- .pollDelay (Duration .ZERO )
645- .pollInterval (Duration .ofMillis (10 ))
646- .untilAsserted (() -> {
647- // We do this in a loop for the wire-side of the connection - data has
648- // to be buffered then flushed
649-
650- // Advance time so the tool emits some data
651- virtualTime .advanceTimeBy (toolEmissionInterval );
652- // Verify the data is received
653- assertThat (responseStream .events ()).anyMatch (line -> line .contains ("notifications/message" ));
654- });
674+ await ().atMost (Duration .ofSeconds (5 )).pollInterval (Duration .ofMillis (20 )).untilTrue (toolCalled );
655675 virtualTime .advanceTimeBy (SESSION_SWEEP_INTERVAL .multipliedBy (2 ));
656676 assertThat (postNotification (httpClient , sessionId )).as ("Session with POST stream is active" )
657677 .isEqualTo (HttpServletResponse .SC_ACCEPTED );
@@ -781,8 +801,8 @@ record StreamResponse(CompletableFuture<Void> streamFuture, Queue<String> events
781801 AtomicReference <InputStream > streamRef ) {
782802
783803 void closeStream () {
784- // Close listening stream. We retry a few times in case the stream was not
785- // established on the first try
804+ // Close the stream. The response headers arrive asynchronously, so wait
805+ // until the client has received them
786806 await ().pollDelay (Duration .ZERO ).atMost (Duration .ofSeconds (1 )).until (() -> {
787807 var stream = streamRef .get ();
788808 if (stream != null ) {
0 commit comments