Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,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;

Expand Down Expand Up @@ -128,6 +133,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.
Expand All @@ -139,25 +150,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
Expand Down Expand Up @@ -195,6 +210,8 @@ public static class Builder {

private SseMessageEndpointValidator messageEndpointValidator = new DefaultSseMessageEndpointValidator();

private int maxResponseSize = DEFAULT_MAX_RESPONSE_SIZE;

/**
* Creates a new builder instance.
*/
Expand Down Expand Up @@ -326,6 +343,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
Expand All @@ -334,7 +371,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);
}

}
Expand All @@ -353,13 +390,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) {
Expand Down Expand Up @@ -490,7 +529,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)));
});
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,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
Expand Down Expand Up @@ -161,11 +166,18 @@ static boolean isMessageEvent(String eventName) {

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,
McpHttpClientTransportAuthorizationErrorHandler authorizationErrorHandler,
List<String> supportedProtocolVersions) {
List<String> supportedProtocolVersions, int maxResponseSize) {
Assert.isTrue(maxResponseSize > 0, "maxResponseSize must be positive");
this.jsonMapper = jsonMapper;
this.httpClient = httpClient;
this.requestBuilder = requestBuilder;
Expand All @@ -181,6 +193,7 @@ private HttpClientStreamableHttpTransport(McpJsonMapper jsonMapper, HttpClient h
.sorted(Comparator.reverseOrder())
.findFirst()
.get();
this.maxResponseSize = maxResponseSize;
}

@Override
Expand Down Expand Up @@ -229,7 +242,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();
}

Expand Down Expand Up @@ -470,16 +484,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;
Expand Down Expand Up @@ -681,6 +695,10 @@ else if (statusCode == BAD_REQUEST) {
return Flux.<McpSchema.JSONRPCMessage>error(new McpTransportException(
"Bad Request. Status code:" + statusCode + ", response-event:" + responseEvent));
}
else if (statusCode >= 400 && statusCode < 500) {
return Flux.<McpSchema.JSONRPCMessage>error(
new McpTransportException("Invalid request. Status code: " + statusCode));
}

return Flux.<McpSchema.JSONRPCMessage>error(
new RuntimeException("Failed to send message: " + responseEvent));
Expand Down Expand Up @@ -750,6 +768,8 @@ public static class Builder {

private McpHttpClientTransportAuthorizationErrorHandler authorizationErrorHandler = McpHttpClientTransportAuthorizationErrorHandler.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
Expand Down Expand Up @@ -948,6 +968,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.
Expand All @@ -957,7 +997,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);
}

}
Expand Down
Loading
Loading