Bound HTTP client reads

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

Author: Kehrlann

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

@@ -82,6 +82,11 @@ public class HttpClientSseClientTransport implements McpClientTransport {
 	/** Default SSE endpoint path */
 	private static final String DEFAULT_SSE_ENDPOINT = "/sse";
 
+	/**
+	 * Default maximum number of bytes read for a single inbound message.
+	 */
+	private static final int DEFAULT_MAX_RESPONSE_SIZE = 16 * 1024 * 1024; // 16MiB
+
 	/** Base URI for the MCP server */
 	private final URI baseUri;
 
@@ -122,6 +127,12 @@ public class HttpClientSseClientTransport implements McpClientTransport {
 	 */
 	private final SseMessageEndpointValidator messageEndpointValidator;
 
+	/**
+	 * Maximum number of bytes read for a single inbound message, whether it arrives on
+	 * the SSE stream or as the response to a posted message.
+	 */
+	private final int maxResponseSize;
+
 	/**
 	 * Creates a new transport instance with custom HTTP client builder, object mapper,
 	 * and headers.
@@ -133,25 +144,29 @@ public class HttpClientSseClientTransport implements McpClientTransport {
 	 * @param httpRequestCustomizer customizer for the requestBuilder before executing
 	 * requests
 	 * @param messageEndpointValidator validator for the message endpoint
+	 * @param maxResponseSize the maximum number of bytes read for a single inbound
+	 * message
 	 * @throws IllegalArgumentException if objectMapper, clientBuilder, or headers is null
 	 */
 	HttpClientSseClientTransport(HttpClient httpClient, HttpRequest.Builder requestBuilder, String baseUri,
 			String sseEndpoint, McpJsonMapper jsonMapper, McpAsyncHttpClientRequestCustomizer httpRequestCustomizer,
-			SseMessageEndpointValidator messageEndpointValidator) {
+			SseMessageEndpointValidator messageEndpointValidator, int maxResponseSize) {
 		Assert.notNull(jsonMapper, "jsonMapper must not be null");
 		Assert.hasText(baseUri, "baseUri must not be empty");
 		Assert.hasText(sseEndpoint, "sseEndpoint must not be empty");
 		Assert.notNull(httpClient, "httpClient must not be null");
 		Assert.notNull(requestBuilder, "requestBuilder must not be null");
 		Assert.notNull(httpRequestCustomizer, "httpRequestCustomizer must not be null");
 		Assert.notNull(messageEndpointValidator, "messageEndpointValidator must not be null");
+		Assert.isTrue(maxResponseSize > 0, "maxResponseSize must be positive");
 		this.baseUri = URI.create(baseUri);
 		this.sseEndpoint = sseEndpoint;
 		this.jsonMapper = jsonMapper;
 		this.httpClient = httpClient;
 		this.requestBuilder = requestBuilder;
 		this.httpRequestCustomizer = httpRequestCustomizer;
 		this.messageEndpointValidator = messageEndpointValidator;
+		this.maxResponseSize = maxResponseSize;
 	}
 
 	@Override
@@ -189,6 +204,8 @@ public static class Builder {
 
 		private SseMessageEndpointValidator messageEndpointValidator = new DefaultSseMessageEndpointValidator();
 
+		private int maxResponseSize = DEFAULT_MAX_RESPONSE_SIZE;
+
 		/**
 		 * Creates a new builder instance.
 		 */
@@ -338,6 +355,26 @@ public Builder messageEndpointValidator(SseMessageEndpointValidator messageEndpo
 			return this;
 		}
 
+		/**
+		 * Sets the maximum number of bytes read for a single inbound message, whether it
+		 * arrives on the SSE stream or as the response to a posted message. A peer that
+		 * sends a larger message (or never terminates one) has its stream aborted instead
+		 * of forcing the transport to buffer it in memory. Defaults to 16MiB.
+		 *
+		 * <p>
+		 * The bound applies per message, not to the stream as a whole: a long-lived SSE
+		 * stream may deliver any number of messages, each up to this size. SSE field
+		 * framing is allowed a small amount of headroom on top of this size, so a message
+		 * of exactly this many bytes is still accepted.
+		 * @param maxResponseSize the maximum inbound message size, in bytes
+		 * @return this builder
+		 */
+		public Builder maxResponseSize(int maxResponseSize) {
+			Assert.isTrue(maxResponseSize > 0, "maxResponseSize must be positive");
+			this.maxResponseSize = maxResponseSize;
+			return this;
+		}
+
 		/**
 		 * Builds a new {@link HttpClientSseClientTransport} instance.
 		 * @return a new transport instance
@@ -346,7 +383,7 @@ public HttpClientSseClientTransport build() {
 			HttpClient httpClient = this.clientBuilder.connectTimeout(this.connectTimeout).build();
 			return new HttpClientSseClientTransport(httpClient, requestBuilder, baseUri, sseEndpoint,
 					jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper, httpRequestCustomizer,
-					messageEndpointValidator);
+					messageEndpointValidator, maxResponseSize);
 		}
 
 	}
@@ -365,13 +402,15 @@ public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> h
 			var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
 			return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, null, transportContext));
 		}).flatMap(requestBuilder -> Mono.create(sink -> {
-			Disposable connection = Flux.<ResponseEvent>create(sseSink -> this.httpClient
-				.sendAsync(requestBuilder.build(),
-						responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink))
-				.exceptionallyCompose(e -> {
-					sseSink.error(e);
-					return CompletableFuture.failedFuture(e);
-				}))
+			Disposable connection = Flux.<ResponseEvent>create(
+					sseSink -> this.httpClient
+						.sendAsync(requestBuilder.build(),
+								responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink,
+										this.maxResponseSize))
+						.exceptionallyCompose(e -> {
+							sseSink.error(e);
+							return CompletableFuture.failedFuture(e);
+						}))
 				.map(responseEvent -> (ResponseSubscribers.SseResponseEvent) responseEvent)
 				.flatMap(responseEvent -> {
 					if (isClosing) {
@@ -502,7 +541,8 @@ 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(httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString()));
+			return Mono.fromFuture(
+					httpClient.sendAsync(request, ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize)));
 		});
 	}
 

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

@@ -8,7 +8,6 @@
 import java.net.URI;
 import java.net.http.HttpClient;
 import java.net.http.HttpRequest;
-import java.net.http.HttpResponse;
 import java.net.http.HttpResponse.BodyHandler;
 import java.time.Duration;
 import java.util.Collections;
@@ -85,6 +84,11 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
 
 	private static final String DEFAULT_ENDPOINT = "/mcp";
 
+	/**
+	 * Default maximum number of bytes read for a single inbound message.
+	 */
+	private static final int DEFAULT_MAX_RESPONSE_SIZE = 16 * 1024 * 1024; // 16MiB
+
 	/**
 	 * HTTP client for sending messages to the server. Uses HTTP POST over the message
 	 * endpoint
@@ -136,10 +140,18 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport {
 
 	private final String latestSupportedProtocolVersion;
 
+	/**
+	 * Maximum number of bytes read for a single inbound message, whether it arrives on an
+	 * SSE stream or as a JSON response body.
+	 */
+	private final int maxResponseSize;
+
 	private HttpClientStreamableHttpTransport(McpJsonMapper jsonMapper, HttpClient httpClient,
 			HttpRequest.Builder requestBuilder, String baseUri, String endpoint, boolean resumableStreams,
 			boolean openConnectionOnStartup, McpAsyncHttpClientRequestCustomizer httpRequestCustomizer,
-			McpHttpClientAuthorizationErrorHandler authorizationErrorHandler, List<String> supportedProtocolVersions) {
+			McpHttpClientAuthorizationErrorHandler authorizationErrorHandler, List<String> supportedProtocolVersions,
+			int maxResponseSize) {
+		Assert.isTrue(maxResponseSize > 0, "maxResponseSize must be positive");
 		this.jsonMapper = jsonMapper;
 		this.httpClient = httpClient;
 		this.requestBuilder = requestBuilder;
@@ -155,6 +167,7 @@ private HttpClientStreamableHttpTransport(McpJsonMapper jsonMapper, HttpClient h
 			.sorted(Comparator.reverseOrder())
 			.findFirst()
 			.get();
+		this.maxResponseSize = maxResponseSize;
 	}
 
 	@Override
@@ -211,7 +224,8 @@ private Publisher<Void> createDelete(String sessionId) {
 			return Mono.from(this.httpRequestCustomizer.customize(builder, "DELETE", uri, null, transportContext));
 		}).flatMap(requestBuilder -> {
 			var request = requestBuilder.build();
-			return Mono.fromFuture(() -> this.httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofString()));
+			return Mono.fromFuture(() -> this.httpClient.sendAsync(request,
+					ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize)));
 		}).then();
 	}
 
@@ -434,16 +448,16 @@ private BodyHandler<Void> toSendMessageBodySubscriber(FluxSink<ResponseEvent> si
 			if (contentType.contains(TEXT_EVENT_STREAM)) {
 				// For SSE streams, use line subscriber that returns Void
 				logger.debug("Received SSE stream response, using line subscriber");
-				return ResponseSubscribers.sseToBodySubscriber(responseInfo, sink);
+				return ResponseSubscribers.sseToBodySubscriber(responseInfo, sink, this.maxResponseSize);
 			}
 			else if (contentType.contains(APPLICATION_JSON)) {
 				// For JSON responses and others, use string subscriber
 				logger.debug("Received response, using string subscriber");
-				return ResponseSubscribers.aggregateBodySubscriber(responseInfo, sink);
+				return ResponseSubscribers.aggregateBodySubscriber(responseInfo, sink, this.maxResponseSize);
 			}
 
 			logger.debug("Received Bodyless response, using discarding subscriber");
-			return ResponseSubscribers.bodilessBodySubscriber(responseInfo, sink);
+			return ResponseSubscribers.bodilessBodySubscriber(responseInfo, sink, this.maxResponseSize);
 		};
 
 		return responseBodyHandler;
@@ -701,6 +715,8 @@ public static class Builder {
 
 		private McpHttpClientAuthorizationErrorHandler authorizationErrorHandler = McpHttpClientAuthorizationErrorHandler.NOOP;
 
+		private int maxResponseSize = DEFAULT_MAX_RESPONSE_SIZE;
+
 		/**
 		 * Creates a new builder with the specified base URI.
 		 * @param baseUri the base URI of the MCP server
@@ -891,6 +907,26 @@ public Builder supportedProtocolVersions(List<String> supportedProtocolVersions)
 			return this;
 		}
 
+		/**
+		 * Sets the maximum number of bytes read for a single inbound message, whether it
+		 * arrives on an SSE stream or as a JSON response body. A peer that sends a larger
+		 * message (or never terminates one) has its stream aborted instead of forcing the
+		 * transport to buffer it in memory. Defaults to 16MiB.
+		 *
+		 * <p>
+		 * The bound applies per message, not to the stream as a whole: a long-lived SSE
+		 * stream may deliver any number of messages, each up to this size. SSE field
+		 * framing is allowed a small amount of headroom on top of this size, so a message
+		 * of exactly this many bytes is still accepted.
+		 * @param maxResponseSize the maximum inbound message size, in bytes
+		 * @return this builder
+		 */
+		public Builder maxResponseSize(int maxResponseSize) {
+			Assert.isTrue(maxResponseSize > 0, "maxResponseSize must be positive");
+			this.maxResponseSize = maxResponseSize;
+			return this;
+		}
+
 		/**
 		 * Construct a fresh instance of {@link HttpClientStreamableHttpTransport} using
 		 * the current builder configuration.
@@ -900,7 +936,7 @@ public HttpClientStreamableHttpTransport build() {
 			HttpClient httpClient = this.clientBuilder.connectTimeout(this.connectTimeout).build();
 			return new HttpClientStreamableHttpTransport(jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper,
 					httpClient, requestBuilder, baseUri, endpoint, resumableStreams, openConnectionOnStartup,
-					httpRequestCustomizer, authorizationErrorHandler, supportedProtocolVersions);
+					httpRequestCustomizer, authorizationErrorHandler, supportedProtocolVersions, maxResponseSize);
 		}
 
 	}