Bound HttpClient inbound reads in the reactive chain

- This discards all BodySubscribers to align with
  a6c88fccf50277ddf705765154baae4caa906e15, which introduced a
  publisher-based architecture instead of using subscribers.
- Bounding the reads now happens within the reactive chains, e.g. in
  Utf8LineDecoder, decodeAggregateResponse and drain* methods.

Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>

Author: Kehrlann

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

@@ -390,23 +390,23 @@ public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> h
 		}).flatMap(requestBuilder -> Mono.create(sink -> {
 			Disposable connection = Mono
 				.fromFuture(() -> this.httpClient.sendAsync(requestBuilder.build(),
-						ResponseSubscribers.boundedPublisherBodyHandler(this.maxResponseSize)))
+						HttpResponse.BodyHandlers.ofPublisher()))
 				.flatMapMany(response -> {
 					if (isClosing) {
 						// 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());
+						return ResponseSubscribers.drain(response.body(), this.maxResponseSize);
 					}
 
 					int statusCode = response.statusCode();
 
 					if (statusCode >= 200 && statusCode < 300) {
-						Flux<String> lines = ResponseSubscribers.decodeLines(response.body());
+						Flux<String> lines = ResponseSubscribers.decodeLines(response.body(), this.maxResponseSize);
 						return ResponseSubscribers.decodeSseResponse(lines, this.maxResponseSize);
 					}
 					else {
-						return ResponseSubscribers.drainThenError(response.body(),
+						return ResponseSubscribers.drainThenError(response.body(), this.maxResponseSize,
 								new RuntimeException("Failed to connect to SSE stream: " + statusCode));
 					}
 				})
@@ -491,17 +491,7 @@ public Mono<Void> sendMessage(JSONRPCMessage message) {
 			}
 
 			return this.serializeMessage(message)
-				.flatMap(body -> sendHttpPost(messageEndpointUri, body).handle((response, sink) -> {
-					if (response.statusCode() != 200 && response.statusCode() != 201 && response.statusCode() != 202
-							&& response.statusCode() != 206) {
-						sink.error(new RuntimeException("Sending message failed with a non-OK HTTP code: "
-								+ response.statusCode() + " - " + response.body()));
-					}
-					else {
-						sink.next(response);
-						sink.complete();
-					}
-				}))
+				.flatMap(body -> sendHttpPost(messageEndpointUri, body))
 				.doOnError(error -> {
 					if (!isClosing) {
 						logger.error("Error sending message: {}", error.getMessage());
@@ -522,7 +512,16 @@ private Mono<String> serializeMessage(final JSONRPCMessage message) {
 		});
 	}
 
-	private Mono<HttpResponse<String>> sendHttpPost(final String endpoint, final String body) {
+	/**
+	 * POSTs {@code body} to {@code endpoint} and consumes the response, failing if the
+	 * server did not accept the message.
+	 *
+	 * <p>
+	 * The response body is streamed rather than aggregated: it is only read as text when
+	 * a non-OK status makes it part of the failure message, and discarded otherwise.
+	 * Either way it has to be consumed, or the connection is never released.
+	 */
+	private Mono<Void> sendHttpPost(final String endpoint, final String body) {
 		final URI requestUri = Utils.resolveUri(baseUri, endpoint);
 		return Mono.deferContextual(ctx -> {
 			var builder = this.requestBuilder.copy()
@@ -534,8 +533,16 @@ private Mono<HttpResponse<String>> sendHttpPost(final String endpoint, final Str
 			return Mono.from(this.httpRequestCustomizer.customize(builder, "POST", requestUri, body, transportContext));
 		}).flatMap(customizedBuilder -> {
 			var request = customizedBuilder.build();
-			return Mono.fromFuture(this.httpClient.sendAsync(request,
-					ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize)));
+			return Mono.fromFuture(this.httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofPublisher()))
+				.flatMap(response -> {
+					int statusCode = response.statusCode();
+					if (statusCode == 200 || statusCode == 201 || statusCode == 202 || statusCode == 206) {
+						return ResponseSubscribers.drain(response.body(), this.maxResponseSize).then();
+					}
+					return ResponseSubscribers.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

@@ -241,8 +241,12 @@ private Publisher<Void> createDelete(String sessionId) {
 			var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
 			return Mono.from(this.httpRequestCustomizer.customize(builder, "DELETE", uri, null, transportContext));
 		})
-			.flatMap(requestBuilder -> Mono.fromFuture(() -> this.httpClient.sendAsync(requestBuilder.build(),
-					ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize))))
+			.flatMap(requestBuilder -> Mono.fromFuture(
+					() -> 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))
+				.then())
 			.then();
 	}
 
@@ -282,7 +286,7 @@ 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);
+		Flux<String> lines = ResponseSubscribers.decodeLines(body, this.maxResponseSize);
 		return ResponseSubscribers.decodeSseResponse(lines, this.maxResponseSize).flatMap(sseEvent -> {
 			if (!isMessageEvent(sseEvent.event())) {
 				logger.debug("Received SSE event with type: {}", sseEvent);
@@ -372,8 +376,7 @@ private Mono<Disposable> reconnect(McpTransportStream<Disposable> stream) {
 				Optional<String> maybeSessionId = request.headers().firstValue(HttpHeaders.MCP_SESSION_ID);
 
 				return Mono
-					.fromFuture(() -> this.httpClient.sendAsync(request,
-							ResponseSubscribers.boundedPublisherBodyHandler(this.maxResponseSize)))
+					.fromFuture(() -> this.httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofPublisher()))
 					.flatMapMany(httpResponse -> {
 						int statusCode = httpResponse.statusCode();
 						Exception exception = null;
@@ -430,8 +433,10 @@ else if (statusCode >= 200 && statusCode < 300) {
 						}
 
 						return proceed ? consumeSseStream(httpResponse.body(), stream, null)
-								: exception != null ? ResponseSubscribers.drainThenError(httpResponse.body(), exception)
-										: ResponseSubscribers.drain(httpResponse.body());
+								: exception != null
+										? ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
+												exception)
+										: ResponseSubscribers.drain(httpResponse.body(), this.maxResponseSize);
 					});
 			})
 				.retryWhen(authorizationErrorRetrySpec())
@@ -540,7 +545,7 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage sentMessage) {
 			})
 				.flatMapMany(requestBuilder -> Mono
 					.fromFuture(() -> this.httpClient.sendAsync(requestBuilder.build(),
-							ResponseSubscribers.boundedPublisherBodyHandler(this.maxResponseSize)))
+							HttpResponse.BodyHandlers.ofPublisher()))
 					.flatMapMany(httpResponse -> {
 						int statusCode = httpResponse.statusCode();
 						Optional<String> maybeSessionId = transportSession == null ? Optional.empty()
@@ -550,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(),
+							return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
 									new McpHttpClientTransportAuthorizationException(
 											"Authorization error when sending message", requestSnapshot,
 											toResponseInfo(httpResponse)));
@@ -575,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());
+								return ResponseSubscribers.drain(httpResponse.body(), this.maxResponseSize);
 							}
 							else if (contentType.contains(TEXT_EVENT_STREAM)) {
 								AtomicBoolean delivered = new AtomicBoolean();
@@ -607,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(),
+							return ResponseSubscribers.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(),
+								return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
 										new McpTransportSessionNotFoundException(
 												"Session not found for session ID: " + sessionRepresentation));
 							}
-							return ResponseSubscribers.drainThenError(httpResponse.body(),
+							return ResponseSubscribers.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(),
+								return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
 										new McpTransportSessionNotFoundException(
 												"Session not found for session ID: " + sessionRepresentation));
 							}
-							return ResponseSubscribers.drainThenError(httpResponse.body(),
+							return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
 									new McpTransportException("Bad Request. Status code:" + statusCode));
 						}
 						else if (statusCode >= 400 && statusCode < 500) {
-							return ResponseSubscribers.drainThenError(httpResponse.body(),
+							return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
 									new McpTransportException("Invalid request. Status code: " + statusCode));
 						}
 
-						return ResponseSubscribers.drainThenError(httpResponse.body(),
+						return ResponseSubscribers.drainThenError(httpResponse.body(), this.maxResponseSize,
 								new RuntimeException("Failed to send message, status code: " + statusCode));
 					})
 					.onErrorMap(CompletionException.class, Throwable::getCause))