Ensure errors are priorities over body limits

Also renamed ResponseSubscribers, refactored test code

Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>

Author: chemicL

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

@@ -396,17 +396,17 @@ public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> h
 						// The body is handed over as a publisher and nothing is read off
 						// the wire until it is subscribed, so it has to be drained even
 						// when its content is of no further interest.
-						return ResponseSubscribers.drain(response.body(), this.maxResponseSize);
+						return ResponseBodyHandlers.drain(response.body(), this.maxResponseSize);
 					}
 
 					int statusCode = response.statusCode();
 
 					if (statusCode >= 200 && statusCode < 300) {
-						Flux<String> lines = ResponseSubscribers.decodeLines(response.body(), this.maxResponseSize);
-						return ResponseSubscribers.decodeSseResponse(lines, this.maxResponseSize);
+						Flux<String> lines = ResponseBodyHandlers.decodeLines(response.body(), this.maxResponseSize);
+						return ResponseBodyHandlers.decodeSseResponse(lines, this.maxResponseSize);
 					}
 					else {
-						return ResponseSubscribers.drainThenError(response.body(), this.maxResponseSize,
+						return ResponseBodyHandlers.drainThenError(response.body(), this.maxResponseSize,
 								new RuntimeException("Failed to connect to SSE stream: " + statusCode));
 					}
 				})
@@ -537,10 +537,10 @@ private Mono<Void> sendHttpPost(final String endpoint, final String body) {
 				.flatMap(response -> {
 					int statusCode = response.statusCode();
 					if (statusCode == 200 || statusCode == 201 || statusCode == 202 || statusCode == 206) {
-						return ResponseSubscribers.drain(response.body(), this.maxResponseSize).then();
+						return ResponseBodyHandlers.drain(response.body(), this.maxResponseSize).then();
 					}
-					return ResponseSubscribers.decodeAggregateResponse(response.body(), this.maxResponseSize)
-						.flatMap(text -> Mono.error(new RuntimeException(
+					return ResponseBodyHandlers.decodeAggregateResponse(response.body(), this.maxResponseSize)
+					                           .flatMap(text -> Mono.error(new RuntimeException(
 								"Sending message failed with a non-OK HTTP code: " + statusCode + " - " + text)));
 				});
 		});

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

@@ -245,7 +245,7 @@ private Publisher<Void> createDelete(String sessionId) {
 					() -> this.httpClient.sendAsync(requestBuilder.build(), HttpResponse.BodyHandlers.ofPublisher()))
 				// The response is not inspected, but the body still has to be consumed
 				// to release the connection.
-				.flatMapMany(response -> ResponseSubscribers.drain(response.body(), this.maxResponseSize))
+				.flatMapMany(response -> ResponseBodyHandlers.drain(response.body(), this.maxResponseSize))
 				.then())
 			.then();
 	}
@@ -286,8 +286,8 @@ public Mono<Void> closeGracefully() {
 	private Flux<McpSchema.JSONRPCMessage> consumeSseStream(
 			java.util.concurrent.Flow.Publisher<List<java.nio.ByteBuffer>> body,
 			McpTransportStream<Disposable> existingStream, Runnable onFirstMessage) {
-		Flux<String> lines = ResponseSubscribers.decodeLines(body, this.maxResponseSize);
-		return ResponseSubscribers.decodeSseResponse(lines, this.maxResponseSize).flatMap(sseEvent -> {
+		Flux<String> lines = ResponseBodyHandlers.decodeLines(body, this.maxResponseSize);
+		return ResponseBodyHandlers.decodeSseResponse(lines, this.maxResponseSize).flatMap(sseEvent -> {
 			if (!isMessageEvent(sseEvent.event())) {
 				logger.debug("Received SSE event with type: {}", sseEvent);
 				if (onFirstMessage != null) {
@@ -434,9 +434,9 @@ else if (statusCode >= 200 && statusCode < 300) {
 
 						return proceed ? consumeSseStream(httpResponse.body(), stream, null)
 								: exception != null
-										? ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+										? ResponseBodyHandlers.drainThenError(httpResponse.body(), this.maxResponseSize,
 												exception)
-										: ResponseSubscribers.drain(httpResponse.body(), this.maxResponseSize);
+										: ResponseBodyHandlers.drain(httpResponse.body(), this.maxResponseSize);
 					});
 			})
 				.retryWhen(authorizationErrorRetrySpec())
@@ -555,7 +555,7 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
 							var request = requestBuilder.build();
 							var requestSnapshot = new HttpRequestSnapshot(request.uri(), request.method(),
 									request.headers());
-							return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+							return ResponseBodyHandlers.drainThenError(httpResponse.body(), this.maxResponseSize,
 									new McpHttpClientTransportAuthorizationException(
 											"Authorization error when sending message", requestSnapshot,
 											toResponseInfo(httpResponse)));
@@ -580,7 +580,7 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
 							if (contentType.isBlank() || "0".equals(contentLength) || statusCode == 202) {
 								logger.debug("No body returned for POST in session {}", sessionRepresentation);
 								deliveredSink.success();
-								return ResponseSubscribers.drain(httpResponse.body(), this.maxResponseSize);
+								return ResponseBodyHandlers.drain(httpResponse.body(), this.maxResponseSize);
 							}
 							else if (contentType.contains(TEXT_EVENT_STREAM)) {
 								AtomicBoolean delivered = new AtomicBoolean();
@@ -591,7 +591,7 @@ else if (contentType.contains(TEXT_EVENT_STREAM)) {
 								});
 							}
 							else if (contentType.contains(APPLICATION_JSON)) {
-								return ResponseSubscribers
+								return ResponseBodyHandlers
 									.decodeAggregateResponse(httpResponse.body(), this.maxResponseSize)
 									.flatMapMany(data -> {
 										deliveredSink.success();
@@ -612,34 +612,34 @@ else if (contentType.contains(APPLICATION_JSON)) {
 
 							logger.warn("Unknown media type {} returned for POST in session {}", contentType,
 									sessionRepresentation);
-							return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+							return ResponseBodyHandlers.drainThenError(httpResponse.body(), this.maxResponseSize,
 									new RuntimeException("Unknown media type returned: " + contentType));
 						}
 						else if (statusCode == NOT_FOUND) {
 							if (maybeSessionId.isPresent()) {
 								logger.debug("Session not found for session ID: {}", sessionRepresentation);
-								return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+								return ResponseBodyHandlers.drainThenError(httpResponse.body(), this.maxResponseSize,
 										new McpTransportSessionNotFoundException(
 												"Session not found for session ID: " + sessionRepresentation));
 							}
-							return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+							return ResponseBodyHandlers.drainThenError(httpResponse.body(), this.maxResponseSize,
 									new McpTransportException("Server Not Found. Status code:" + statusCode));
 						}
 						else if (statusCode == BAD_REQUEST) {
 							if (maybeSessionId.isPresent()) {
-								return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+								return ResponseBodyHandlers.drainThenError(httpResponse.body(), this.maxResponseSize,
 										new McpTransportSessionNotFoundException(
 												"Session not found for session ID: " + sessionRepresentation));
 							}
-							return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+							return ResponseBodyHandlers.drainThenError(httpResponse.body(), this.maxResponseSize,
 									new McpTransportException("Bad Request. Status code:" + statusCode));
 						}
 						else if (statusCode >= 400 && statusCode < 500) {
-							return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+							return ResponseBodyHandlers.drainThenError(httpResponse.body(), this.maxResponseSize,
 									new McpTransportException("Invalid request. Status code: " + statusCode));
 						}
 
-						return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+						return ResponseBodyHandlers.drainThenError(httpResponse.body(), this.maxResponseSize,
 								new RuntimeException("Failed to send message, status code: " + statusCode));
 					})
 					.onErrorMap(CompletionException.class, Throwable::getCause))

…lient/transport/ResponseSubscribers.java …ient/transport/ResponseBodyHandlers.javamcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java renamed to mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseBodyHandlers.java mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java renamed to mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseBodyHandlers.java

@@ -36,7 +36,7 @@
  * @author Dariusz Jędrzejczyk
  * @author Daniel Garnier-Moiroux
  */
-class ResponseSubscribers {
+class ResponseBodyHandlers {
 
 	/**
 	 * Bytes of SSE field framing a single line may carry on top of the message payload:
@@ -137,7 +137,7 @@ static Mono<String> decodeAggregateResponse(Publisher<List<ByteBuffer>> publishe
 	 * @param error the error to propagate once the body has been discarded
 	 */
 	static <T> Flux<T> drainThenError(Publisher<List<ByteBuffer>> body, int maxSize, Throwable error) {
-		return boundTotalBytes(body, maxSize).thenMany(Flux.error(error));
+		return boundTotalBytes(body, maxSize).onErrorComplete().thenMany(Mono.error(error));
 	}
 
 	/**

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

@@ -17,7 +17,7 @@
 import reactor.adapter.JdkFlowAdapter;
 import reactor.core.publisher.Flux;
 
-import io.modelcontextprotocol.client.transport.ResponseSubscribers.SseEvent;
+import io.modelcontextprotocol.client.transport.ResponseBodyHandlers.SseEvent;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -48,7 +48,7 @@
  *
  * <p>
  * The middle column read each chunk incrementally, which is ~25x quicker than what 2.0.0
- * shipped, but {@link ResponseSubscribers.Utf8LineDecoder} still searched its buffered
+ * shipped, but {@link ResponseBodyHandlers.Utf8LineDecoder} still searched its buffered
  * characters for a line terminator from the start of the buffer on every chunk, so eight
  * times the payload cost ~45x the time. Resuming that search where the previous one ended
  * gives the third column, which scales with the payload rather than with its square and
@@ -173,8 +173,8 @@ private static long timeDecode(byte[] body, int expectedEvents) {
 	private static List<SseEvent> decode(byte[] body) {
 		Flow.Publisher<List<ByteBuffer>> publisher = JdkFlowAdapter
 			.publisherToFlowPublisher(Flux.fromIterable(chunk(body)));
-		Flux<String> lines = ResponseSubscribers.decodeLines(publisher, Integer.MAX_VALUE);
-		return ResponseSubscribers.decodeSseResponse(lines, MAX_SIZE).collectList().block();
+		Flux<String> lines = ResponseBodyHandlers.decodeLines(publisher, Integer.MAX_VALUE);
+		return ResponseBodyHandlers.decodeSseResponse(lines, MAX_SIZE).collectList().block();
 	}
 
 	private static List<List<ByteBuffer>> chunk(byte[] body) {

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

@@ -8,8 +8,8 @@
 
 import org.junit.jupiter.api.Test;
 
-import io.modelcontextprotocol.client.transport.ResponseSubscribers.SseEvent;
-import io.modelcontextprotocol.client.transport.ResponseSubscribers.SseEventParser;
+import io.modelcontextprotocol.client.transport.ResponseBodyHandlers.SseEvent;
+import io.modelcontextprotocol.client.transport.ResponseBodyHandlers.SseEventParser;
 
 import static org.assertj.core.api.Assertions.assertThat;
 

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

@@ -8,7 +8,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 
-import io.modelcontextprotocol.client.transport.ResponseSubscribers.Utf8LineDecoder;
+import io.modelcontextprotocol.client.transport.ResponseBodyHandlers.Utf8LineDecoder;
 import io.modelcontextprotocol.spec.McpTransportException;
 import org.junit.jupiter.api.Test;
 

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

@@ -9,7 +9,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 
-import io.modelcontextprotocol.client.transport.ResponseSubscribers.Utf8LineDecoder;
+import io.modelcontextprotocol.client.transport.ResponseBodyHandlers.Utf8LineDecoder;
 import org.junit.jupiter.api.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;

mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientBoundedReadTestSupport.java

@@ -8,14 +8,18 @@
 import java.io.OutputStream;
 import java.net.InetSocketAddress;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.Arrays;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executors;
 
-import com.sun.net.httpserver.HttpExchange;
 import com.sun.net.httpserver.HttpServer;
 import io.modelcontextprotocol.server.transport.TomcatTestUtil;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 
+import static org.assertj.core.api.Assertions.assertThat;
+
 /**
  * Shared fixture for the transport bounded-read tests: a bare {@link HttpServer} whose
  * response body is written by a per-test {@link Responder}.
@@ -62,67 +66,79 @@ void stopServer() {
 	}
 
 	/**
-	 * Registers a handler that responds with the given content type and body.
+	 * Registers a handler that answers {@code method} requests to {@code path} with the
+	 * given content type and a body written by {@code responder}. Any other method gets a
+	 * 405, as from a server offering nothing else there, so that requests the test does
+	 * not target, such as the GET stream the Streamable HTTP transport opens once
+	 * initialized, neither reach the responder nor stand in for the targeted request.
+	 * @return completes once the targeted response has been handled: with the
+	 * {@link IOException} that cut the body short if the client hung up first, or with
+	 * {@code null} if the body was written in full
 	 */
-	protected void respondWith(String path, String contentType, Responder responder) {
+	protected CompletableFuture<IOException> respondWith(String method, String path, String contentType,
+			Responder responder) {
+		CompletableFuture<IOException> response = new CompletableFuture<>();
 		this.server.createContext(path, exchange -> {
-			exchange.getResponseHeaders().set("Content-Type", contentType);
-			exchange.sendResponseHeaders(200, 0);
-			try (OutputStream body = exchange.getResponseBody()) {
-				responder.respond(body);
-			}
-			catch (IOException ignored) {
-				// The client aborts the response once the limit is exceeded, which closes
-				// the connection and makes further writes fail. That is the behaviour
-				// under test.
+			try {
+				if (!method.equals(exchange.getRequestMethod())) {
+					exchange.sendResponseHeaders(405, -1);
+					return;
+				}
+				exchange.getResponseHeaders().set("Content-Type", contentType);
+				exchange.sendResponseHeaders(200, 0);
+				try (OutputStream body = exchange.getResponseBody()) {
+					responder.respond(body);
+					response.complete(null);
+				}
+				catch (IOException ex) {
+					response.complete(ex);
+				}
 			}
 			finally {
 				exchange.close();
 			}
 		});
+		return response;
 	}
 
 	/**
-	 * A responder that writes {@code chunks} blocks of {@code 'a'} with no line
-	 * terminator anywhere, so nothing downstream can ever flush a line.
+	 * Asserts that the client hung up on {@code response}, which is how exceeding the
+	 * bound must end: with the endless responders below, a client that read the body in
+	 * full would never let it complete, and one that merely stopped reading would leave
+	 * the server blocked writing into a stalled connection.
 	 */
-	protected static Responder unterminatedLine(int chunks) {
-		return body -> {
-			byte[] chunk = new byte[MAX_SIZE];
-			java.util.Arrays.fill(chunk, (byte) 'a');
-			for (int i = 0; i < chunks; i++) {
-				body.write(chunk);
-				body.flush();
-			}
-		};
+	protected static void assertHungUp(CompletableFuture<IOException> response) {
+		assertThat(response).succeedsWithin(Duration.ofSeconds(5)).isNotNull();
 	}
 
 	/**
-	 * A responder that writes {@code chunks} blocks of {@code 'a'}, each ending in a lone
-	 * CR. The line decoder only flushes a line on LF, so its buffer keeps growing even
-	 * though a CR arrives regularly.
+	 * A responder that streams {@code 'a'} with no line terminator anywhere, so nothing
+	 * downstream can ever flush a line.
 	 */
-	protected static Responder carriageReturnTerminatedRuns(int chunks) {
-		return body -> {
-			byte[] chunk = new byte[MAX_SIZE];
-			java.util.Arrays.fill(chunk, (byte) 'a');
-			chunk[MAX_SIZE - 1] = '\r';
-			for (int i = 0; i < chunks; i++) {
-				body.write(chunk);
-				body.flush();
-			}
-		};
+	protected static Responder unterminatedLine() {
+		byte[] block = new byte[MAX_SIZE];
+		Arrays.fill(block, (byte) 'a');
+		return endlessly(block);
 	}
 
 	/**
-	 * A responder that writes enough short, properly terminated lines to exceed the limit
-	 * in aggregate.
+	 * A responder that streams short, properly terminated lines, each small but exceeding
+	 * the limit in aggregate.
 	 */
 	protected static Responder manyShortLines(String prefix) {
+		return endlessly((prefix + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n").getBytes(StandardCharsets.UTF_8));
+	}
+
+	/**
+	 * A responder that repeats {@code block} until the client hangs up, as a peer
+	 * streaming without end would. There is no amount to tune: however much the socket
+	 * buffers absorb, and whether the client closes with a FIN or a RST, the writes only
+	 * stop once the connection is gone.
+	 */
+	private static Responder endlessly(byte[] block) {
 		return body -> {
-			byte[] line = (prefix + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\n").getBytes(StandardCharsets.UTF_8);
-			for (int i = 0; i < (MAX_SIZE / line.length) + 64; i++) {
-				body.write(line);
+			while (true) {
+				body.write(block);
 				body.flush();
 			}
 		};

mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransportBoundedReadTests.java

@@ -48,38 +48,41 @@ void releaseStream() {
 	void shouldRejectSingleLineExceedingMaxSize() {
 		// A line that never terminates, so the line buffer underneath the SSE parser
 		// would grow without limit before any event could be flushed.
-		respondWith(endpoint(), "text/event-stream", unterminatedLine(8));
+		var response = respondWith("GET", endpoint(), "text/event-stream", unterminatedLine());
 
 		StepVerifier.create(connect())
 			.verifyErrorMatches(t -> messageContains(t, "Inbound line exceeds the maximum allowed size"));
+		assertHungUp(response);
 	}
 
 	@Test
 	void shouldRejectEventExceedingMaxSize() {
 		// Many short, terminated "data:" lines with no blank line to end the event. Each
 		// line is small, but the accumulated event data would grow without limit.
-		respondWith(endpoint(), "text/event-stream", manyShortLines("data:"));
+		var response = respondWith("GET", endpoint(), "text/event-stream", manyShortLines("data:"));
 
 		StepVerifier.create(connect())
 			.verifyErrorMatches(t -> messageContains(t, "Inbound SSE event exceeds the maximum allowed size"));
+		assertHungUp(response);
 	}
 
 	@Test
-	void shouldRejectPostResponseExceedingMaxSize() throws Exception {
+	void shouldRejectPostResponseExceedingMaxSize() {
 		// The response to a posted message is discarded on success, but a peer must still
 		// not be able to make the transport read an unbounded one.
-		respondWith(endpoint(), "text/event-stream", body -> {
+		respondWith("GET", endpoint(), "text/event-stream", body -> {
 			body.write(("event:endpoint\ndata:" + MESSAGE_ENDPOINT + "\n\n").getBytes(StandardCharsets.UTF_8));
 			body.flush();
 			awaitTeardown();
 		});
-		respondWith(MESSAGE_ENDPOINT, "application/json", unterminatedLine(64));
+		var response = respondWith("POST", MESSAGE_ENDPOINT, "application/json", unterminatedLine());
 
 		HttpClientSseClientTransport transport = transport();
 		transport.connect(Function.identity()).block(Duration.ofSeconds(5));
 
 		StepVerifier.create(sendMessage(transport))
 			.verifyErrorMatches(t -> messageContains(t, "Inbound response body exceeds the maximum allowed size"));
+		assertHungUp(response);
 	}
 
 	private void awaitTeardown() {