Skip to content

Commit 5eef75a

Browse files
chemicLchemicL
committed
Handle Streamable transport similarly with switchOnNext
Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
1 parent 9796c7c commit 5eef75a

3 files changed

Lines changed: 199 additions & 71 deletions

File tree

‎mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java‎

Lines changed: 69 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
import java.util.Optional;
1818
import java.util.concurrent.CompletionException;
1919
import java.util.concurrent.Flow;
20-
import java.util.concurrent.atomic.AtomicBoolean;
2120
import java.util.concurrent.atomic.AtomicReference;
2221
import java.util.function.Consumer;
2322
import java.util.function.Function;
@@ -269,6 +268,15 @@ private void handleException(Throwable t) {
269268
}
270269
}
271270

271+
private void handleExceptionSafely(Throwable t) {
272+
try {
273+
handleException(t);
274+
}
275+
catch (Exception e) {
276+
logger.error("Error handling exception {}", t.getMessage(), e);
277+
}
278+
}
279+
272280
@Override
273281
public Mono<Void> closeGracefully() {
274282
return Mono.defer(() -> {
@@ -282,38 +290,33 @@ public Mono<Void> closeGracefully() {
282290
});
283291
}
284292

285-
private Flux<McpSchema.JSONRPCMessage> consumeSseStream(Flow.Publisher<List<ByteBuffer>> body,
286-
McpTransportStream<Disposable> existingStream, Runnable onFirstMessage) {
293+
/**
294+
* Every successfully processed event yields exactly one element, empty when it
295+
* carries no message, so that callers can tell when the first one has arrived.
296+
*/
297+
private Flux<Optional<McpSchema.JSONRPCMessage>> consumeSseStream(Flow.Publisher<List<ByteBuffer>> body,
298+
McpTransportStream<Disposable> existingStream) {
287299
Flux<String> lines = ResponseBodyHandlers.decodeLines(body, this.maxResponseSize);
288300
return ResponseBodyHandlers.decodeSseResponse(lines, this.maxResponseSize).flatMap(sseEvent -> {
289301
if (!isMessageEvent(sseEvent.event())) {
290302
logger.debug("Received SSE event with type: {}", sseEvent);
291-
if (onFirstMessage != null) {
292-
onFirstMessage.run();
293-
}
294-
return Flux.empty();
303+
return Flux.just(Optional.empty());
295304
}
296305
String data = sseEvent.data();
297306
if (data == null || data.isBlank()) {
298307
logger.debug("Skipping SSE event with empty data (stream primer)");
299-
if (onFirstMessage != null) {
300-
onFirstMessage.run();
301-
}
302-
return Flux.empty();
308+
return Flux.just(Optional.empty());
303309
}
304310
try {
305311
McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.jsonMapper, data);
306312
Tuple2<Optional<String>, Iterable<McpSchema.JSONRPCMessage>> idWithMessages = Tuples
307313
.of(Optional.ofNullable(sseEvent.id()), List.of(message));
308314
McpTransportStream<Disposable> sessionStream = existingStream != null ? existingStream
309315
: new DefaultMcpTransportStream<>(this.resumableStreams, this::reconnect);
310-
if (onFirstMessage != null) {
311-
onFirstMessage.run();
312-
}
313-
return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages)));
316+
return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages))).map(Optional::of);
314317
}
315318
catch (IOException e) {
316-
return Flux.<McpSchema.JSONRPCMessage>error(
319+
return Flux.<Optional<McpSchema.JSONRPCMessage>>error(
317320
new McpTransportException("Error parsing JSON-RPC message: " + data, e));
318321
}
319322
});
@@ -395,10 +398,11 @@ private Mono<Disposable> reconnect(McpTransportStream<Disposable> stream) {
395398
"Unrecognized server error when connecting to SSE stream, status code: " + statusCode);
396399
}
397400
logger.debug("SSE connection established successfully");
398-
return consumeSseStream(httpResponse.body(), stream, null);
401+
return consumeSseStream(httpResponse.body(), stream);
399402
});
400403
})
401-
.retryWhen(authorizationErrorRetrySpec())
404+
.retryWhen(authorizationErrorRetrySpec()).<McpSchema
405+
.JSONRPCMessage>handle((message, messages) -> message.ifPresent(messages::next))
402406
.flatMap(jsonrpcMessage -> requestHandler.apply(Mono.just(jsonrpcMessage)))
403407
.onErrorComplete(t -> {
404408
if (t instanceof CompletionException) {
@@ -480,15 +484,6 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
480484

481485
final AtomicReference<Disposable> disposableRef = new AtomicReference<>();
482486

483-
// Once sendMessage() has completed, a later failure can no longer be reported
484-
// through its sink: signalling it there would only have Reactor drop it.
485-
final AtomicBoolean delivered = new AtomicBoolean();
486-
final Runnable markDelivered = () -> {
487-
if (delivered.compareAndSet(false, true)) {
488-
deliveredSink.success();
489-
}
490-
};
491-
492487
Disposable connection = Mono.deferContextual(ctx -> {
493488
HttpRequest.Builder requestBuilder = this.requestBuilder.copy();
494489

@@ -543,29 +538,32 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
543538

544539
if (contentType.isBlank() || "0".equals(contentLength) || statusCode == 202) {
545540
logger.debug("No body returned for POST in session {}", sessionRepresentation);
546-
markDelivered.run();
547-
return ResponseBodyHandlers.drain(httpResponse.body(), this.maxResponseSize);
541+
return ResponseBodyHandlers.<Optional<McpSchema.JSONRPCMessage>>drain(httpResponse.body(),
542+
this.maxResponseSize)
543+
.startWith(Optional.empty());
548544
}
549545
else if (contentType.contains(TEXT_EVENT_STREAM)) {
550-
return consumeSseStream(httpResponse.body(), null, markDelivered);
546+
return consumeSseStream(httpResponse.body(), null);
551547
}
552548
else if (contentType.contains(APPLICATION_JSON)) {
553-
return ResponseBodyHandlers.decodeAggregateResponse(httpResponse.body(), this.maxResponseSize)
554-
.flatMapMany(data -> {
555-
markDelivered.run();
556-
if (sentMessage instanceof McpSchema.JSONRPCNotification) {
557-
logger.warn("Notification: {} received non-compliant response: {}", sentMessage,
558-
Utils.hasText(data) ? data : "[empty]");
559-
return Flux.empty();
560-
}
561-
try {
562-
return Flux.just(McpSchema.deserializeJsonRpcMessage(jsonMapper, data));
563-
}
564-
catch (IOException e) {
565-
return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
566-
"Error deserializing JSON-RPC message: " + data, e));
567-
}
568-
});
549+
return ResponseBodyHandlers.decodeAggregateResponse(httpResponse.body(),
550+
this.maxResponseSize).<Optional<McpSchema.JSONRPCMessage>>handle((data, messages) -> {
551+
if (sentMessage instanceof McpSchema.JSONRPCNotification) {
552+
logger.warn("Notification: {} received non-compliant response: {}", sentMessage,
553+
Utils.hasText(data) ? data : "[empty]");
554+
messages.next(Optional.empty());
555+
return;
556+
}
557+
try {
558+
messages
559+
.next(Optional.of(McpSchema.deserializeJsonRpcMessage(jsonMapper, data)));
560+
}
561+
catch (IOException e) {
562+
messages.error(new McpTransportException(
563+
"Error deserializing JSON-RPC message: " + data, e));
564+
}
565+
})
566+
.flux();
569567
}
570568

571569
logger.warn("Unknown media type {} returned for POST in session {}", contentType,
@@ -575,32 +573,41 @@ else if (contentType.contains(APPLICATION_JSON)) {
575573
});
576574
})
577575
.retryWhen(authorizationErrorRetrySpec())
578-
.flatMap(jsonRpcMessage -> requestHandler.apply(Mono.just(jsonRpcMessage)))
579576
.onErrorMap(CompletionException.class, t -> t.getCause())
577+
// sendMessage() is resolved by the first signal only: any later failure
578+
// is
579+
// merely handled below, as sendMessage() has already completed by then.
580+
// An exchange ending without any event still means the server accepted
581+
// the message, so completion resolves it successfully too.
582+
.switchOnFirst((first, messages) -> {
583+
if (first.isOnError()) {
584+
// Handled before failing sendMessage(), so that a session the
585+
// server does not recognise is already invalidated by the time
586+
// the caller learns about it. Consumed here so that it is not
587+
// handled a second time below.
588+
handleExceptionSafely(first.getThrowable());
589+
deliveredSink.error(first.getThrowable());
590+
return Flux.empty();
591+
}
592+
deliveredSink.success();
593+
return messages;
594+
}).<McpSchema
595+
.JSONRPCMessage>handle((message, messages) -> message.ifPresent(messages::next))
596+
.flatMap(jsonRpcMessage -> requestHandler.apply(Mono.just(jsonRpcMessage)))
580597
.doFinally(s -> {
581598
Disposable ref = disposableRef.getAndSet(null);
582599
if (ref != null) {
583600
transportSession.removeConnection(ref);
584601
}
585602
})
586603
.onErrorComplete(t -> {
587-
// handle the error first
588-
try {
589-
this.handleException(t);
590-
}
591-
catch (Exception e) {
592-
logger.error("Error handling exception {}", t.getMessage(), e);
593-
}
594-
// inform the caller of sendMessage, unless it has already completed
595-
if (delivered.compareAndSet(false, true)) {
596-
deliveredSink.error(t);
597-
}
604+
handleExceptionSafely(t);
598605
return true;
599606
})
600-
// An exchange can end without anything having signalled delivery, e.g.
601-
// an SSE response closed before its first event. The server accepted
602-
// the message all the same, so sendMessage() must not be left pending.
603-
.doOnComplete(markDelivered)
607+
// Closing the session before the first signal cancels the exchange:
608+
// complete sendMessage() instead of leaving it pending. A no-op once it
609+
// has resolved.
610+
.doOnCancel(deliveredSink::success)
604611
.contextWrite(deliveredSink.contextView())
605612
.subscribe();
606613

‎mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportSendMessageTests.java‎

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
/*
2+
* Copyright 2026-2026 the original author or authors.
3+
*/
4+
5+
package io.modelcontextprotocol.client.transport;
6+
7+
import java.io.IOException;
8+
import java.io.OutputStream;
9+
import java.net.InetSocketAddress;
10+
import java.nio.charset.StandardCharsets;
11+
import java.time.Duration;
12+
import java.util.concurrent.CountDownLatch;
13+
import java.util.concurrent.ExecutorService;
14+
import java.util.concurrent.Executors;
15+
import java.util.concurrent.TimeUnit;
16+
17+
import com.sun.net.httpserver.HttpExchange;
18+
import com.sun.net.httpserver.HttpHandler;
19+
import com.sun.net.httpserver.HttpServer;
20+
import io.modelcontextprotocol.spec.McpSchema;
21+
import io.modelcontextprotocol.spec.McpTransportException;
22+
import io.modelcontextprotocol.spec.json.gson.GsonMcpJsonMapper;
23+
import org.junit.jupiter.api.AfterEach;
24+
import org.junit.jupiter.api.Test;
25+
import reactor.test.StepVerifier;
26+
27+
import static org.assertj.core.api.Assertions.assertThat;
28+
29+
/**
30+
* Verifies that {@link HttpClientStreamableHttpTransport#sendMessage} always resolves,
31+
* and that it fails when the server's response to it cannot be read.
32+
*/
33+
class HttpClientStreamableHttpTransportSendMessageTests {
34+
35+
// Only bounds a regression: every test resolves without waiting on it.
36+
private static final Duration TIMEOUT = Duration.ofSeconds(5);
37+
38+
private static final McpSchema.JSONRPCRequest REQUEST = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION,
39+
"ping", "1", null);
40+
41+
private final ExecutorService executor = Executors.newCachedThreadPool();
42+
43+
private final CountDownLatch releaseResponse = new CountDownLatch(1);
44+
45+
private HttpServer server;
46+
47+
@AfterEach
48+
void tearDown() {
49+
this.releaseResponse.countDown();
50+
if (this.server != null) {
51+
this.server.stop(0);
52+
}
53+
this.executor.shutdownNow();
54+
}
55+
56+
@Test
57+
void sendMessageFailsWhenJsonResponseIsMalformed() throws IOException {
58+
HttpClientStreamableHttpTransport transport = transport(exchange -> {
59+
byte[] body = "{broken".getBytes(StandardCharsets.UTF_8);
60+
exchange.getResponseHeaders().add("Content-Type", "application/json");
61+
exchange.sendResponseHeaders(200, body.length);
62+
try (OutputStream outputStream = exchange.getResponseBody()) {
63+
outputStream.write(body);
64+
}
65+
});
66+
67+
StepVerifier.create(transport.sendMessage(REQUEST))
68+
.expectErrorSatisfies(e -> assertThat(e).isInstanceOf(McpTransportException.class)
69+
.hasMessageContaining("Error deserializing JSON-RPC message"))
70+
.verify(TIMEOUT);
71+
}
72+
73+
@Test
74+
void sendMessageCompletesWhenClosedBeforeAnyEvent() throws IOException {
75+
CountDownLatch streamOpened = new CountDownLatch(1);
76+
HttpClientStreamableHttpTransport transport = transport(exchange -> {
77+
exchange.getResponseHeaders().add("Content-Type", "text/event-stream");
78+
exchange.sendResponseHeaders(200, 0);
79+
exchange.getResponseBody().flush();
80+
streamOpened.countDown();
81+
try {
82+
this.releaseResponse.await();
83+
}
84+
catch (InterruptedException e) {
85+
Thread.currentThread().interrupt();
86+
}
87+
exchange.close();
88+
});
89+
90+
StepVerifier.create(transport.sendMessage(REQUEST)).then(() -> {
91+
try {
92+
assertThat(streamOpened.await(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)).isTrue();
93+
}
94+
catch (InterruptedException e) {
95+
throw new IllegalStateException(e);
96+
}
97+
transport.closeGracefully().block(TIMEOUT);
98+
}).expectComplete().verify(TIMEOUT);
99+
}
100+
101+
private HttpClientStreamableHttpTransport transport(HttpHandler postHandler) throws IOException {
102+
this.server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
103+
this.server.setExecutor(this.executor);
104+
this.server.createContext("/mcp", exchange -> {
105+
if ("POST".equals(exchange.getRequestMethod())) {
106+
postHandler.handle(exchange);
107+
}
108+
else {
109+
// No standalone SSE stream, which keeps the POST the only exchange.
110+
methodNotAllowed(exchange);
111+
}
112+
});
113+
this.server.start();
114+
return HttpClientStreamableHttpTransport.builder("http://127.0.0.1:" + this.server.getAddress().getPort())
115+
.jsonMapper(new GsonMcpJsonMapper())
116+
.build();
117+
}
118+
119+
private static void methodNotAllowed(HttpExchange exchange) throws IOException {
120+
try (exchange) {
121+
exchange.sendResponseHeaders(405, -1);
122+
}
123+
}
124+
125+
}

‎mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportEmptyResponseTests.java‎

Lines changed: 5 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,6 @@
2828
import io.modelcontextprotocol.client.transport.customizer.McpSyncHttpClientRequestCustomizer;
2929
import io.modelcontextprotocol.server.transport.TomcatTestUtil;
3030
import io.modelcontextprotocol.spec.McpSchema;
31-
import io.modelcontextprotocol.spec.ProtocolVersions;
3231
import reactor.test.StepVerifier;
3332

3433
/**
@@ -112,18 +111,15 @@ void testNotificationInitialized() throws URISyntaxException {
112111
.httpRequestCustomizer(mockRequestCustomizer)
113112
.build();
114113

115-
var initializeRequest = McpSchema.InitializeRequest
116-
.builder(ProtocolVersions.MCP_2025_03_26, McpSchema.ClientCapabilities.builder().roots(true).build(),
117-
McpSchema.Implementation.builder("MCP Client", "0.3.1").build())
118-
.build();
119-
var testMessage = new McpSchema.JSONRPCRequest(McpSchema.METHOD_INITIALIZE, "test-id", initializeRequest);
114+
// Some servers answer a notification with an empty JSON body rather than 202.
115+
var testMessage = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
116+
McpSchema.METHOD_NOTIFICATION_INITIALIZED, null);
120117

121118
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
122119

123120
// Verify the customizer was called
124-
verify(mockRequestCustomizer, atLeastOnce()).customize(any(), eq("POST"), eq(uri), eq(
125-
"{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":\"test-id\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{\"roots\":{\"listChanged\":true}},\"clientInfo\":{\"name\":\"MCP Client\",\"version\":\"0.3.1\"}}}"),
126-
any());
121+
verify(mockRequestCustomizer, atLeastOnce()).customize(any(), eq("POST"), eq(uri),
122+
eq("{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}"), any());
127123

128124
}
129125

0 commit comments

Comments
 (0)