From bf4760324707745b217c7243cd6b59d7099e4199 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dariusz=20J=C4=99drzejczyk?= <2554306+chemicL@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:53:09 +0200 Subject: [PATCH 1/2] Refactor HttpClient use to avoid capturing the sseSink that leads to leaking the httpClient reference due to a cycle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Daniel Garnier-Moiroux Signed-off-by: Dariusz Jędrzejczyk <2554306+chemicL@users.noreply.github.com> --- .../HttpClientSseClientTransport.java | 88 ++- .../HttpClientStreamableHttpTransport.java | 503 ++++++-------- .../client/transport/ResponseSubscribers.java | 626 ++++++++---------- .../spec/DefaultMcpTransportSession.java | 1 + .../transport/BoundedBodySubscriberTests.java | 50 +- .../transport/HttpClientLeakTestSupport.java | 58 ++ ...HttpClientSseClientTransportLeakTests.java | 51 ++ ...lientStreamableHttpTransportLeakTests.java | 55 ++ .../transport/LoopbackMcpHttpServer.java | 136 ++++ .../client/transport/SseEventParserTests.java | 120 ++++ .../transport/Utf8LineDecoderTests.java | 135 ++++ .../spec/DefaultMcpTransportSessionTests.java | 26 + .../HttpClientBoundedReadTestSupport.java | 17 + ...reamableHttpTransportBoundedReadTests.java | 10 + 14 files changed, 1180 insertions(+), 696 deletions(-) create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientLeakTestSupport.java create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransportLeakTests.java create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportLeakTests.java create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/transport/LoopbackMcpHttpServer.java create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/transport/SseEventParserTests.java create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/spec/DefaultMcpTransportSessionTests.java diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java index 3e4b613ff..64badceec 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java @@ -11,12 +11,10 @@ import java.net.http.HttpResponse; import java.time.Duration; import java.util.List; -import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.Function; -import io.modelcontextprotocol.client.transport.ResponseSubscribers.ResponseEvent; import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientRequestCustomizer; import io.modelcontextprotocol.client.transport.customizer.McpSyncHttpClientRequestCustomizer; import io.modelcontextprotocol.common.McpTransportContext; @@ -390,61 +388,59 @@ public Mono connect(Function, Mono> 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.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 -> { + Disposable connection = Mono + .fromFuture(() -> this.httpClient.sendAsync(requestBuilder.build(), + ResponseSubscribers.boundedPublisherBodyHandler(this.maxResponseSize))) + .flatMapMany(response -> { if (isClosing) { - return Mono.empty(); + return Flux.empty(); } - int statusCode = responseEvent.responseInfo().statusCode(); + int statusCode = response.statusCode(); if (statusCode >= 200 && statusCode < 300) { - try { - if (ENDPOINT_EVENT_TYPE.equals(responseEvent.sseEvent().event())) { - String messageEndpointUri = responseEvent.sseEvent().data(); - try { - messageEndpointValidator.validate(uri, messageEndpointUri); - } - catch (InvalidSseMessageEndpointException e) { - sink.error(e); - this.messageEndpointSink.tryEmitError(e); - return Flux.error(e); - } - if (this.messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) { - sink.success(); - return Flux.empty(); // No further processing needed - } - else { - sink.error(new RuntimeException("Failed to handle SSE endpoint event")); - } + Flux lines = ResponseSubscribers.decodeLines(response.body()); + return ResponseSubscribers.decodeSseResponse(lines, this.maxResponseSize); + } + else { + return ResponseSubscribers.drainThenError(response.body(), + new RuntimeException("Failed to connect to SSE stream: " + statusCode)); + } + }) + .flatMap(sseEvent -> { + try { + if (ENDPOINT_EVENT_TYPE.equals(sseEvent.event())) { + String messageEndpointUri = sseEvent.data(); + try { + messageEndpointValidator.validate(uri, messageEndpointUri); + } + catch (InvalidSseMessageEndpointException e) { + sink.error(e); + this.messageEndpointSink.tryEmitError(e); + return Flux.error(e); } - else if (MESSAGE_EVENT_TYPE.equals(responseEvent.sseEvent().event())) { - JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(jsonMapper, - responseEvent.sseEvent().data()); + if (this.messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) { sink.success(); - return Flux.just(message); + return Flux.empty(); // No further processing needed } else { - logger.debug("Received unrecognized SSE event type: {}", responseEvent.sseEvent()); - sink.success(); + sink.error(new RuntimeException("Failed to handle SSE endpoint event")); } } - catch (IOException e) { - sink.error(new McpTransportException("Error processing SSE event", e)); + else if (MESSAGE_EVENT_TYPE.equals(sseEvent.event())) { + JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(jsonMapper, sseEvent.data()); + sink.success(); + return Flux.just(message); + } + else { + logger.debug("Received unrecognized SSE event type: {}", sseEvent); + sink.success(); } } - return Flux.error( - new RuntimeException("Failed to send message: " + responseEvent)); - + catch (IOException e) { + sink.error(new McpTransportException("Error processing SSE event", e)); + } + return Flux.empty(); }) .flatMap(jsonRpcMessage -> handler.apply(Mono.just(jsonRpcMessage))) .onErrorComplete(t -> { @@ -529,8 +525,8 @@ private Mono> 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, ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize))); + return Mono.fromFuture(this.httpClient.sendAsync(request, + ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize))); }); } diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java index 07a8f5e23..8d518af35 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java @@ -9,19 +9,20 @@ import java.net.http.HttpClient; import java.net.http.HttpRequest; import java.net.http.HttpResponse; -import java.net.http.HttpResponse.BodyHandler; +import java.nio.ByteBuffer; import java.time.Duration; import java.util.Collections; import java.util.Comparator; import java.util.List; import java.util.Optional; import java.util.concurrent.CompletionException; +import java.util.concurrent.Flow; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.Function; import io.modelcontextprotocol.client.McpAsyncClient; -import io.modelcontextprotocol.client.transport.ResponseSubscribers.ResponseEvent; import io.modelcontextprotocol.client.transport.customizer.McpAsyncHttpClientRequestCustomizer; import io.modelcontextprotocol.client.transport.customizer.McpHttpClientAuthorizationErrorHandler; import io.modelcontextprotocol.client.transport.customizer.McpHttpClientTransportAuthorizationErrorHandler; @@ -49,7 +50,6 @@ import org.slf4j.LoggerFactory; import reactor.core.Disposable; import reactor.core.publisher.Flux; -import reactor.core.publisher.FluxSink; import reactor.core.publisher.Mono; import reactor.util.function.Tuple2; import reactor.util.function.Tuples; @@ -207,7 +207,7 @@ public static Builder builder(String baseUri) { @Override public Mono connect(Function, Mono> handler) { - return Mono.deferContextual(ctx -> { + return Mono.defer(() -> { this.handler.set(handler); if (this.openConnectionOnStartup) { logger.debug("Eagerly opening connection on startup"); @@ -240,11 +240,10 @@ private Publisher createDelete(String sessionId) { .DELETE(); var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY); return Mono.from(this.httpRequestCustomizer.customize(builder, "DELETE", uri, null, transportContext)); - }).flatMap(requestBuilder -> { - var request = requestBuilder.build(); - return Mono.fromFuture(() -> this.httpClient.sendAsync(request, - ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize))); - }).then(); + }) + .flatMap(requestBuilder -> Mono.fromFuture(() -> this.httpClient.sendAsync(requestBuilder.build(), + ResponseSubscribers.boundedStringBodyHandler(this.maxResponseSize)))) + .then(); } @Override @@ -254,7 +253,8 @@ public void setExceptionHandler(Consumer handler) { } private void handleException(Throwable t) { - logger.debug("Handling exception for session {}", sessionIdOrPlaceholder(this.activeSession.get()), t); + logger.debug("Handling exception for session {}", sessionIdOrPlaceholder( + activeSession.get() != null ? activeSession.get().sessionId() : Optional.empty()), t); if (t instanceof McpTransportSessionNotFoundException) { McpTransportSession invalidSession = this.activeSession.getAndSet(createTransportSession()); logger.warn("Server does not recognize session {}. Invalidating.", invalidSession.sessionId()); @@ -279,6 +279,41 @@ public Mono closeGracefully() { }); } + private Flux consumeSseStream( + java.util.concurrent.Flow.Publisher> body, + McpTransportStream existingStream, Runnable onFirstMessage) { + Flux lines = ResponseSubscribers.decodeLines(body); + return ResponseSubscribers.decodeSseResponse(lines, this.maxResponseSize).flatMap(sseEvent -> { + if (!isMessageEvent(sseEvent.event())) { + logger.debug("Received SSE event with type: {}", sseEvent); + return Flux.empty(); + } + String data = sseEvent.data(); + if (data == null || data.isBlank()) { + logger.debug("Skipping SSE event with empty data (stream primer)"); + if (onFirstMessage != null) { + onFirstMessage.run(); + } + return Flux.empty(); + } + try { + McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.jsonMapper, data); + Tuple2, Iterable> idWithMessages = Tuples + .of(Optional.ofNullable(sseEvent.id()), List.of(message)); + McpTransportStream sessionStream = existingStream != null ? existingStream + : new DefaultMcpTransportStream<>(this.resumableStreams, this::reconnect); + if (onFirstMessage != null) { + onFirstMessage.run(); + } + return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages))); + } + catch (IOException e) { + return Flux.error( + new McpTransportException("Error parsing JSON-RPC message", e)); + } + }); + } + private Mono reconnect(McpTransportStream stream) { return Mono.deferContextual(ctx -> { var rh = this.handler.get(); @@ -304,14 +339,15 @@ private Mono reconnect(McpTransportStream stream) { final AtomicReference disposableRef = new AtomicReference<>(); - var uri = Utils.resolveUri(this.baseUri, this.endpoint); + Optional maybeSessionId = transportSession == null ? Optional.empty() + : transportSession.sessionId(); Disposable connection = Mono.deferContextual(connectionCtx -> { + var uri = Utils.resolveUri(this.baseUri, this.endpoint); HttpRequest.Builder requestBuilder = this.requestBuilder.copy(); - if (transportSession != null && transportSession.sessionId().isPresent()) { - requestBuilder = requestBuilder.header(HttpHeaders.MCP_SESSION_ID, - transportSession.sessionId().get()); + if (maybeSessionId.isPresent()) { + requestBuilder = requestBuilder.header(HttpHeaders.MCP_SESSION_ID, maybeSessionId.get()); } if (stream != null && stream.lastId().isPresent()) { @@ -328,123 +364,84 @@ private Mono reconnect(McpTransportStream stream) { var transportContext = connectionCtx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY); return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, null, transportContext)); }) - .flatMapMany(requestBuilder -> Flux.create(sseSink -> this.httpClient - .sendAsync(requestBuilder.build(), this.toSendMessageBodySubscriber(sseSink)) - .whenComplete((response, throwable) -> { - if (throwable != null) { - sseSink.error(throwable); - } - else { - logger.debug("SSE connection established successfully"); - } - })).flatMap(responseEvent -> { - int statusCode = responseEvent.responseInfo().statusCode(); + .flatMapMany(requestBuilder -> Mono + .fromFuture(() -> this.httpClient.sendAsync(requestBuilder.build(), + ResponseSubscribers.boundedPublisherBodyHandler(this.maxResponseSize))) + .flatMapMany(httpResponse -> { + int statusCode = httpResponse.statusCode(); + Exception exception = null; + boolean proceed = false; if (statusCode == 401 || statusCode == 403) { logger.debug("Authorization error in reconnect with code {}", statusCode); var request = requestBuilder.build(); var requestSnapshot = new HttpRequestSnapshot(request.uri(), request.method(), request.headers()); - return Mono.error( - new McpHttpClientTransportAuthorizationException( - "Authorization error connecting to SSE stream", requestSnapshot, - responseEvent.responseInfo())); + exception = new McpHttpClientTransportAuthorizationException( + "Authorization error connecting to SSE stream", requestSnapshot, + toResponseInfo(httpResponse)); } else if (statusCode == METHOD_NOT_ALLOWED) { logger.debug("The server does not support SSE streams, using request-response mode."); - return Flux.empty(); - } - - if (!(responseEvent instanceof ResponseSubscribers.SseResponseEvent sseResponseEvent)) { - return Flux.error(new McpTransportException( - "Unrecognized server error when connecting to SSE stream, status code: " - + statusCode)); } - else if (statusCode >= 200 && statusCode < 300) { - if (isMessageEvent(sseResponseEvent.sseEvent().event())) { - String data = sseResponseEvent.sseEvent().data(); - // Per 2025-11-25 spec (SEP-1699), servers may - // send SSE events - // with empty data to prime the client for - // reconnection. - // Skip these events as they contain no JSON-RPC - // message. - if (data == null || data.isBlank()) { - logger.debug("Skipping SSE event with empty data (stream primer)"); - return Flux.empty(); - } - try { - // We don't support batching ATM and probably - // won't since the next version considers - // removing it. - McpSchema.JSONRPCMessage message = McpSchema - .deserializeJsonRpcMessage(this.jsonMapper, data); - - Tuple2, Iterable> idWithMessages = Tuples - .of(Optional.ofNullable(sseResponseEvent.sseEvent().id()), List.of(message)); - - McpTransportStream sessionStream = stream != null ? stream - : new DefaultMcpTransportStream<>(this.resumableStreams, this::reconnect); - logger.debug("Connected stream {}", sessionStream.streamId()); - - return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages))); - - } - catch (IOException ioException) { - return Flux.error(new McpTransportException( - "Error parsing JSON-RPC message: " + responseEvent, ioException)); - } + else if (statusCode == NOT_FOUND) { + if (maybeSessionId.isPresent()) { + logger.debug("Session not found for session ID: {}", maybeSessionId.get()); + String sessionIdRepresentation = sessionIdOrPlaceholder(maybeSessionId); + exception = new McpTransportSessionNotFoundException(sessionIdRepresentation); } else { - logger.debug("Received SSE event with type: {}", sseResponseEvent.sseEvent()); - return Flux.empty(); + exception = new McpTransportException("Server Not Found. Status code:" + statusCode); } } - else if (statusCode == NOT_FOUND) { - - if (transportSession != null && transportSession.sessionId().isPresent()) { - // only if the request was sent with a session id - // and the response is 404, we consider it a - // session not found error. - logger.debug("Session not found for session ID: {}", - transportSession.sessionId().get()); - String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession); - McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException( - "Session not found for session ID: " + sessionIdRepresentation); - return Flux.error(exception); + else if (statusCode == BAD_REQUEST) { + // if the session id was set, but the session no longer + // exists, some servers can return 400 instead of 404 + if (maybeSessionId.isPresent()) { + String sessionIdRepresentation = sessionIdOrPlaceholder(maybeSessionId); + exception = new McpTransportSessionNotFoundException(sessionIdRepresentation); + } + else { + exception = new McpTransportException("Bad Request. Status code:" + statusCode); } - return Flux.error( - new McpTransportException("Server Not Found. Status code:" + statusCode - + ", response-event:" + responseEvent)); } - else if (statusCode == BAD_REQUEST) { - if (transportSession != null && transportSession.sessionId().isPresent()) { - // only if the request was sent with a session id - // and thre response is 404, we consider it a - // session not found error. - String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession); - McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException( - "Session not found for session ID: " + sessionIdRepresentation); - return Flux.error(exception); + else if (statusCode >= 200 && statusCode < 300) { + String contentType = httpResponse.headers() + .firstValue(HttpHeaders.CONTENT_TYPE) + .orElse("") + .toLowerCase(); + if (contentType.contains(TEXT_EVENT_STREAM)) { + logger.debug("SSE connection established successfully"); + proceed = true; + } + else { + exception = new McpTransportException( + "Unrecognized server error when connecting to SSE stream, status code: " + + statusCode); } - return Flux.error(new McpTransportException( - "Bad Request. Status code:" + statusCode + ", response-event:" + responseEvent)); } - return Flux.error(new McpTransportException( - "Received unrecognized SSE event type: " + sseResponseEvent.sseEvent().event())); - }) - .retryWhen(authorizationErrorRetrySpec()) - .flatMap(jsonrpcMessage -> requestHandler.apply(Mono.just(jsonrpcMessage))) - .onErrorMap(CompletionException.class, t -> t.getCause()) - .doFinally(s -> { - Disposable ref = disposableRef.getAndSet(null); - if (ref != null) { - transportSession.removeConnection(ref); + else { + exception = new McpTransportException("Received unrecognized status code: " + statusCode); } + + return proceed ? consumeSseStream(httpResponse.body(), stream, null) + : exception != null ? ResponseSubscribers.drainThenError(httpResponse.body(), exception) + : ResponseSubscribers.drain(httpResponse.body()); })) + .retryWhen(authorizationErrorRetrySpec()) + .flatMap(jsonrpcMessage -> requestHandler.apply(Mono.just(jsonrpcMessage))) .onErrorComplete(t -> { + if (t instanceof CompletionException) { + t = t.getCause(); + } this.handleException(t); return true; }) + .doFinally(s -> { + Disposable ref = disposableRef.getAndSet(null); + if (ref != null) { + transportSession.removeConnection(ref); + } + }) .contextWrite(ctx) .subscribe(); @@ -455,6 +452,14 @@ else if (statusCode == BAD_REQUEST) { } + private static HttpResponse.ResponseInfo toResponseInfo(HttpResponse>> response) { + return new HttpClientResponseInfo(response.statusCode(), response.headers(), response.version()); + } + + private record HttpClientResponseInfo(int statusCode, java.net.http.HttpHeaders headers, + HttpClient.Version version) implements HttpResponse.ResponseInfo { + } + private Retry authorizationErrorRetrySpec() { return Retry.from(companion -> companion.flatMap(retrySignal -> { if (!(retrySignal.failure() instanceof McpHttpClientTransportAuthorizationException authException)) { @@ -475,31 +480,6 @@ private Retry authorizationErrorRetrySpec() { })); } - private BodyHandler toSendMessageBodySubscriber(FluxSink sink) { - - BodyHandler responseBodyHandler = responseInfo -> { - - String contentType = responseInfo.headers().firstValue(HttpHeaders.CONTENT_TYPE).orElse("").toLowerCase(); - - 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, 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, this.maxResponseSize); - } - - logger.debug("Received Bodyless response, using discarding subscriber"); - return ResponseSubscribers.bodilessBodySubscriber(responseInfo, sink, this.maxResponseSize); - }; - - return responseBodyHandler; - - } - public String toString(McpSchema.JSONRPCMessage message) { try { return this.jsonMapper.writeValueAsString(message); @@ -529,9 +509,6 @@ public Mono sendMessage(McpSchema.JSONRPCMessage sentMessage) { final AtomicReference disposableRef = new AtomicReference<>(); - var uri = Utils.resolveUri(this.baseUri, this.endpoint); - String jsonBody = this.toString(sentMessage); - Disposable connection = Mono.deferContextual(ctx -> { HttpRequest.Builder requestBuilder = this.requestBuilder.copy(); @@ -540,6 +517,8 @@ public Mono sendMessage(McpSchema.JSONRPCMessage sentMessage) { transportSession.sessionId().get()); } + String jsonBody = this.toString(sentMessage); + var uri = Utils.resolveUri(this.baseUri, this.endpoint); var builder = requestBuilder.uri(uri) .header(HttpHeaders.ACCEPT, APPLICATION_JSON + ", " + TEXT_EVENT_STREAM) .header(HttpHeaders.CONTENT_TYPE, APPLICATION_JSON_UTF8) @@ -551,168 +530,120 @@ public Mono sendMessage(McpSchema.JSONRPCMessage sentMessage) { var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY); return Mono .from(this.httpRequestCustomizer.customize(builder, "POST", uri, jsonBody, transportContext)); - }).flatMapMany(requestBuilder -> Flux.create(responseEventSink -> { - // Create the async request with proper body subscriber selection - Mono.fromFuture(this.httpClient - .sendAsync(requestBuilder.build(), this.toSendMessageBodySubscriber(responseEventSink)) - .whenComplete((response, throwable) -> { - if (throwable != null) { - responseEventSink.error(throwable); - } - else { - logger.debug("SSE connection established successfully"); + }) + .flatMapMany(requestBuilder -> Mono + .fromFuture(() -> this.httpClient.sendAsync(requestBuilder.build(), + ResponseSubscribers.boundedPublisherBodyHandler(this.maxResponseSize))) + .flatMapMany(httpResponse -> { + int statusCode = httpResponse.statusCode(); + Exception exception = null; + boolean proceed = false; + Optional maybeSessionId = transportSession == null ? Optional.empty() + : transportSession.sessionId(); + if (statusCode == 401 || statusCode == 403) { + logger.debug("Authorization error in sendMessage with code {}", statusCode); + var request = requestBuilder.build(); + var requestSnapshot = new HttpRequestSnapshot(request.uri(), request.method(), + request.headers()); + return ResponseSubscribers.drainThenError(httpResponse.body(), + new McpHttpClientTransportAuthorizationException( + "Authorization error when sending message", requestSnapshot, + toResponseInfo(httpResponse))); } - })).onErrorMap(CompletionException.class, t -> t.getCause()).onErrorComplete().subscribe(); - - }).flatMap(responseEvent -> { - int statusCode = responseEvent.responseInfo().statusCode(); - if (statusCode == 401 || statusCode == 403) { - var request = requestBuilder.build(); - var requestSnapshot = new HttpRequestSnapshot(request.uri(), request.method(), request.headers()); - logger.debug("Authorization error in sendMessage with code {}", statusCode); - return Mono.error(new McpHttpClientTransportAuthorizationException( - "Authorization error when sending message", requestSnapshot, responseEvent.responseInfo())); - } - if (transportSession.markInitialized( - responseEvent.responseInfo().headers().firstValue("mcp-session-id").orElseGet(() -> null))) { - // Once we have a session, we try to open an async stream for - // the server to send notifications and requests out-of-band. + if (transportSession + .markInitialized(httpResponse.headers().firstValue("mcp-session-id").orElse(null)) + && !openConnectionOnStartup) { + reconnect(null).contextWrite(deliveredSink.contextView()).subscribe(); + } - reconnect(null).contextWrite(deliveredSink.contextView()).subscribe(); - } + String sessionRepresentation = sessionIdOrPlaceholder(maybeSessionId); + + if (statusCode >= 200 && statusCode < 300) { + String contentType = httpResponse.headers() + .firstValue(HttpHeaders.CONTENT_TYPE) + .orElse("") + .toLowerCase(); + String contentLength = httpResponse.headers() + .firstValue(HttpHeaders.CONTENT_LENGTH) + .orElse(null); + + 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()); + } + else if (contentType.contains(TEXT_EVENT_STREAM)) { + AtomicBoolean delivered = new AtomicBoolean(); + return consumeSseStream(httpResponse.body(), null, () -> { + if (delivered.compareAndSet(false, true)) { + deliveredSink.success(); + } + }); + } + else if (contentType.contains(APPLICATION_JSON)) { + return ResponseSubscribers + .decodeAggregateResponse(httpResponse.body(), this.maxResponseSize) + .flatMapMany(data -> { + deliveredSink.success(); + if (sentMessage instanceof McpSchema.JSONRPCNotification) { + logger.warn("Notification: {} received non-compliant response: {}", + sentMessage, Utils.hasText(data) ? data : "[empty]"); + return Flux.empty(); + } + try { + return Flux.just(McpSchema.deserializeJsonRpcMessage(jsonMapper, data)); + } + catch (IOException e) { + return Flux.error(new McpTransportException( + "Error deserializing JSON-RPC message", e)); + } + }); + } - String sessionRepresentation = sessionIdOrPlaceholder(transportSession); - - if (statusCode >= 200 && statusCode < 300) { - - String contentType = responseEvent.responseInfo() - .headers() - .firstValue(HttpHeaders.CONTENT_TYPE) - .orElse("") - .toLowerCase(); - - String contentLength = responseEvent.responseInfo() - .headers() - .firstValue(HttpHeaders.CONTENT_LENGTH) - .orElse(null); - - // For empty content or HTTP code 202 (ACCEPTED), assume success - if (contentType.isBlank() || "0".equals(contentLength) || statusCode == 202) { - // if (contentType.isBlank() || "0".equals(contentLength)) { - logger.debug("No body returned for POST in session {}", sessionRepresentation); - // No content type means no response body, so we can just - // return an empty stream - deliveredSink.success(); - return Flux.empty(); - } - else if (contentType.contains(TEXT_EVENT_STREAM)) { - return Flux.just(((ResponseSubscribers.SseResponseEvent) responseEvent).sseEvent()) - .flatMap(sseEvent -> { - String data = sseEvent.data(); - // Per 2025-11-25 spec (SEP-1699), servers may send SSE - // events - // with empty data to prime the client for reconnection. - // Skip these events as they contain no JSON-RPC message. - if (data == null || data.isBlank()) { - logger.debug("Skipping SSE event with empty data (stream primer)"); - return Flux.empty(); - } - try { - // We don't support batching ATM and probably - // won't - // since the - // next version considers removing it. - McpSchema.JSONRPCMessage message = McpSchema - .deserializeJsonRpcMessage(this.jsonMapper, data); - - Tuple2, Iterable> idWithMessages = Tuples - .of(Optional.ofNullable(sseEvent.id()), List.of(message)); - - McpTransportStream sessionStream = new DefaultMcpTransportStream<>( - this.resumableStreams, this::reconnect); - - logger.debug("Connected stream {}", sessionStream.streamId()); - - deliveredSink.success(); - - return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages))); - } - catch (IOException ioException) { - return Flux.error(new McpTransportException( - "Error parsing JSON-RPC message: " + responseEvent, ioException)); - } - }); - } - else if (contentType.contains(APPLICATION_JSON)) { - deliveredSink.success(); - String data = ((ResponseSubscribers.AggregateResponseEvent) responseEvent).data(); - if (sentMessage instanceof McpSchema.JSONRPCNotification) { - logger.warn("Notification: {} received non-compliant response: {}", sentMessage, - Utils.hasText(data) ? data : "[empty]"); - return Mono.empty(); + logger.warn("Unknown media type {} returned for POST in session {}", contentType, + sessionRepresentation); + return ResponseSubscribers.drainThenError(httpResponse.body(), + new RuntimeException("Unknown media type returned: " + contentType)); } - - try { - return Mono.just(McpSchema.deserializeJsonRpcMessage(jsonMapper, data)); + else if (statusCode == NOT_FOUND) { + if (maybeSessionId.isPresent()) { + logger.debug("Session not found for session ID: {}", sessionRepresentation); + return ResponseSubscribers.drainThenError(httpResponse.body(), + new McpTransportSessionNotFoundException( + "Session not found for session ID: " + sessionRepresentation)); + } + return ResponseSubscribers.drainThenError(httpResponse.body(), + new McpTransportException("Server Not Found. Status code:" + statusCode)); } - catch (IOException e) { - return Mono.error(new McpTransportException( - "Error deserializing JSON-RPC message: " + responseEvent, e)); + else if (statusCode == BAD_REQUEST) { + if (maybeSessionId.isPresent()) { + return ResponseSubscribers.drainThenError(httpResponse.body(), + new McpTransportSessionNotFoundException( + "Session not found for session ID: " + sessionRepresentation)); + } + return ResponseSubscribers.drainThenError(httpResponse.body(), + new McpTransportException("Bad Request. Status code:" + statusCode)); + } + else if (statusCode >= 400 && statusCode < 500) { + return ResponseSubscribers.drainThenError(httpResponse.body(), + new McpTransportException("Invalid request. Status code: " + statusCode)); } - } - logger.warn("Unknown media type {} returned for POST in session {}", contentType, - sessionRepresentation); - - return Flux.error( - new RuntimeException("Unknown media type returned: " + contentType)); - } - else if (statusCode == NOT_FOUND) { - if (transportSession != null && transportSession.sessionId().isPresent()) { - // only if the request was sent with a session id and the - // response is 404, we consider it a session not found error. - logger.debug("Session not found for session ID: {}", transportSession.sessionId().get()); - McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException( - "Session not found for session ID: " + sessionRepresentation); - return Flux.error(exception); - } - return Flux.error(new McpTransportException( - "Server Not Found. Status code:" + statusCode + ", response-event:" + responseEvent)); - } - else if (statusCode == BAD_REQUEST) { - // Some implementations can return 400 when presented with a - // session id that it doesn't know about, so we will - // invalidate the session - // https://github.com/modelcontextprotocol/typescript-sdk/issues/389 - - if (transportSession != null && transportSession.sessionId().isPresent()) { - // only if the request was sent with a session id and the - // response is 404, we consider it a session not found error. - McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException( - "Session not found for session ID: " + sessionRepresentation); - return Flux.error(exception); - } - return Flux.error(new McpTransportException( - "Bad Request. Status code:" + statusCode + ", response-event:" + responseEvent)); - } - else if (statusCode >= 400 && statusCode < 500) { - return Flux.error( - new McpTransportException("Invalid request. Status code: " + statusCode)); - } - return Flux.error( - new RuntimeException("Failed to send message: " + responseEvent)); - }) + return ResponseSubscribers.drainThenError(httpResponse.body(), + new RuntimeException("Failed to send message, status code: " + statusCode)); + }) + .onErrorMap(CompletionException.class, Throwable::getCause)) .retryWhen(authorizationErrorRetrySpec()) .flatMap(jsonRpcMessage -> requestHandler.apply(Mono.just(jsonRpcMessage))) .onErrorMap(CompletionException.class, t -> t.getCause()) .doFinally(s -> { - logger.debug("SendMessage finally: {}", s); Disposable ref = disposableRef.getAndSet(null); if (ref != null) { transportSession.removeConnection(ref); } - })).onErrorComplete(t -> { + }) + .onErrorComplete(t -> { // handle the error first try { this.handleException(t); @@ -723,7 +654,9 @@ else if (statusCode >= 400 && statusCode < 500) { // inform the caller of sendMessage deliveredSink.error(t); return true; - }).contextWrite(deliveredSink.contextView()).subscribe(); + }) + .contextWrite(deliveredSink.contextView()) + .subscribe(); disposableRef.set(connection); transportSession.addConnection(connection); @@ -731,8 +664,8 @@ else if (statusCode >= 400 && statusCode < 500) { } - private static String sessionIdOrPlaceholder(McpTransportSession transportSession) { - return transportSession.sessionId().orElse("[missing_session_id]"); + private static String sessionIdOrPlaceholder(Optional sessionId) { + return sessionId.orElse("[missing_session_id]"); } @Override diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java index b19904de6..202f525ca 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java @@ -1,37 +1,40 @@ /* -* Copyright 2024 - 2024 the original author or authors. -*/ + * Copyright 2024 - 2026 the original author or authors. + */ package io.modelcontextprotocol.client.transport; import java.net.http.HttpResponse; import java.net.http.HttpResponse.BodyHandler; import java.net.http.HttpResponse.BodySubscriber; -import java.net.http.HttpResponse.ResponseInfo; import java.nio.ByteBuffer; +import java.nio.CharBuffer; +import java.nio.charset.CharacterCodingException; +import java.nio.charset.CharsetDecoder; +import java.nio.charset.CoderResult; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.concurrent.CompletionStage; import java.util.concurrent.Flow; -import java.util.concurrent.atomic.AtomicReference; -import java.util.regex.Pattern; +import java.util.concurrent.Flow.Publisher; -import org.reactivestreams.FlowAdapters; -import org.reactivestreams.Subscription; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import io.modelcontextprotocol.spec.McpTransportException; -import reactor.core.publisher.BaseSubscriber; -import reactor.core.publisher.FluxSink; +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Flux; +import reactor.core.publisher.Mono; /** - * Utility class providing various {@link BodySubscriber} implementations for handling - * different types of HTTP response bodies in the context of Model Context Protocol (MCP) - * clients. + * Utility class providing various operations for handling different types of HTTP + * response bodies in the context of Model Context Protocol (MCP) clients. * *

- * Defines subscribers for processing Server-Sent Events (SSE), aggregate responses, and - * bodiless responses. + * Defines Flux operators for processing Server-Sent Events (SSE), aggregate responses, + * and bodiless responses. * * @author Christian Tzolov * @author Dariusz Jędrzejczyk @@ -39,8 +42,6 @@ */ class ResponseSubscribers { - private static final Logger logger = LoggerFactory.getLogger(ResponseSubscribers.class); - /** * Bytes of SSE field framing a single line may carry on top of the message payload: * {@code "event: "} is the longest field prefix this parser recognises. Line @@ -53,39 +54,43 @@ class ResponseSubscribers { record SseEvent(String id, String event, String data) { } - sealed interface ResponseEvent permits SseResponseEvent, AggregateResponseEvent, DummyEvent { - - ResponseInfo responseInfo(); - - } - - record DummyEvent(ResponseInfo responseInfo) implements ResponseEvent { - - } - - record SseResponseEvent(ResponseInfo responseInfo, SseEvent sseEvent) implements ResponseEvent { - } - - record AggregateResponseEvent(ResponseInfo responseInfo, String data) implements ResponseEvent { + /** + * Creates a {@link BodyHandler} that exposes the response body as a publisher of + * byte-buffer chunks, bounding how much memory reading a single line may occupy. + * + *

+ * The line decoder downstream of this handler buffers characters until it encounters + * a line terminator, so a peer that never terminates a line (or sends an enormous + * one) would force the transport to buffer it in memory. The bound is allowed + * {@link #SSE_FRAMING_OVERHEAD} extra bytes so that the SSE framing around a payload + * does not count against the payload's own budget; the content type is not known when + * the bound is installed, so the same headroom applies to non-SSE bodies. + * + *

+ * This only bounds a single line. What accumulates across lines is bounded where it + * accumulates: see {@link #decodeSseResponse} for multi-line SSE events and + * {@link #decodeAggregateResponse} for whole response bodies. + * @param maxSize the maximum number of bytes read for a single inbound message + */ + static BodyHandler>> boundedPublisherBodyHandler(int maxSize) { + BodyHandler>> delegate = HttpResponse.BodyHandlers.ofPublisher(); + int bound = plusFramingOverhead(maxSize); + return responseInfo -> new BoundedLineBodySubscriber<>(delegate.apply(responseInfo), bound); } /** - * Creates a {@link BodySubscriber} that parses a Server-Sent Events stream, bounding - * how much memory a single inbound message may occupy. Both the size of an individual - * line (as read off the wire before a terminator is seen) and the accumulated size of - * a multi-line SSE event are capped at {@code maxSize}; a peer exceeding either limit - * has its stream aborted instead of forcing the transport to buffer it in memory. The - * line bound is allowed {@link #SSE_FRAMING_OVERHEAD} extra bytes so that the SSE - * framing around a payload does not count against the payload's own budget. - * @param responseInfo the HTTP response information - * @param sink the sink to emit parsed events to - * @param maxSize the maximum number of bytes read for a single inbound message + * Creates a {@link BodyHandler} that reads the response body into a string, bounding + * how much memory it may occupy. A peer sending more than {@code maxSize} bytes has + * its response aborted instead of forcing the transport to buffer it in memory. + * + *

+ * Decoding matches {@link HttpResponse.BodyHandlers#ofString()}, including its + * handling of the charset declared in the {@code Content-Type} header. + * @param maxSize the maximum number of bytes read for the response body */ - static BodySubscriber sseToBodySubscriber(ResponseInfo responseInfo, FluxSink sink, - int maxSize) { - BodySubscriber lineSubscriber = HttpResponse.BodySubscribers - .fromLineSubscriber(FlowAdapters.toFlowSubscriber(new SseLineSubscriber(responseInfo, sink, maxSize))); - return new BoundedLineBodySubscriber(lineSubscriber, plusFramingOverhead(maxSize)); + static BodyHandler boundedStringBodyHandler(int maxSize) { + BodyHandler delegate = HttpResponse.BodyHandlers.ofString(); + return responseInfo -> new BoundedTotalBodySubscriber<>(delegate.apply(responseInfo), maxSize); } /** @@ -98,347 +103,271 @@ private static int plusFramingOverhead(int maxSize) { } /** - * Creates a {@link BodySubscriber} that aggregates the whole response body into a - * single event, bounding how much memory it may occupy. Both the size of an - * individual line (as read off the wire before a terminator is seen) and the total - * accumulated body are capped at {@code maxSize}; a peer exceeding either limit has - * its response aborted instead of forcing the transport to buffer it in memory. - * @param responseInfo the HTTP response information - * @param sink the sink to emit the aggregated event to - * @param maxSize the maximum number of bytes read for the response body + * Converts a publisher of byte-buffer chunks into a flux of decoded string lines. */ - static BodySubscriber aggregateBodySubscriber(ResponseInfo responseInfo, FluxSink sink, - int maxSize) { - BodySubscriber lineSubscriber = HttpResponse.BodySubscribers - .fromLineSubscriber(FlowAdapters.toFlowSubscriber(new AggregateSubscriber(responseInfo, sink, maxSize))); - return new BoundedLineBodySubscriber(lineSubscriber, maxSize); + static Flux decodeLines(Publisher> publisher) { + return Flux.defer(() -> { + Utf8LineDecoder dec = new Utf8LineDecoder(); + return JdkFlowAdapter.flowPublisherToFlux(publisher) + .concatMapIterable(dec::decode) + .concatWith(Flux.defer(() -> Flux.fromIterable(dec.flush()))); + }); } /** - * Creates a {@link BodySubscriber} that discards the response body, bounding how much - * memory reading it may occupy. The body is discarded as it arrives, but the - * underlying line subscriber still buffers each line before handing it over, so a - * peer sending a line longer than {@code maxSize} has its response aborted. - * @param responseInfo the HTTP response information - * @param sink the sink to emit the completion event to - * @param maxSize the maximum number of bytes read for a single line + * Parses a flux of SSE-formatted lines into a flux of {@link SseEvent}, bounding how + * much memory a single event may occupy. + * @param lines the SSE-formatted lines to parse + * @param maxSize the maximum number of bytes that may accumulate for a single SSE + * event */ - static BodySubscriber bodilessBodySubscriber(ResponseInfo responseInfo, FluxSink sink, - int maxSize) { - BodySubscriber lineSubscriber = HttpResponse.BodySubscribers - .fromLineSubscriber(FlowAdapters.toFlowSubscriber(new BodilessResponseLineSubscriber(responseInfo, sink))); - return new BoundedLineBodySubscriber(lineSubscriber, maxSize); + static Flux decodeSseResponse(Flux lines, int maxSize) { + return Flux.defer(() -> { + SseEventParser parser = new SseEventParser(maxSize); + return lines.handle((line, sink) -> parser.feed(line).ifPresent(sink::next)) + .concatWith(Mono.defer(() -> parser.flush().map(Mono::just).orElseGet(Mono::empty))); + }); } /** - * Creates a {@link BodyHandler} that reads the response body into a string, bounding - * how much memory it may occupy. A peer sending more than {@code maxSize} bytes has - * its response aborted instead of forcing the transport to buffer it in memory. - * - *

- * Decoding matches {@link HttpResponse.BodyHandlers#ofString()}, including its - * handling of the charset declared in the {@code Content-Type} header. + * Collects all byte-buffer chunks from the publisher into a single UTF-8 decoded + * string, bounding how much memory it may occupy. A peer sending a body larger than + * {@code maxSize} has its response aborted instead of forcing the transport to buffer + * it in memory. + * @param publisher the response body * @param maxSize the maximum number of bytes read for the response body */ - static BodyHandler boundedStringBodyHandler(int maxSize) { - BodyHandler delegate = HttpResponse.BodyHandlers.ofString(); - return responseInfo -> new BoundedTotalBodySubscriber<>(delegate.apply(responseInfo), maxSize); + static Mono decodeAggregateResponse(Publisher> publisher, int maxSize) { + return Flux.defer(() -> { + // Held in an array because the handle callback below cannot mutate a + // captured local. The enclosing defer gives each subscriber its own. + long[] totalBytes = new long[1]; + return JdkFlowAdapter.flowPublisherToFlux(publisher) + .flatMapIterable(list -> list) + .handle((buffer, sink) -> { + totalBytes[0] += buffer.remaining(); + if (totalBytes[0] > maxSize) { + sink.error(new McpTransportException( + "Inbound response body exceeds the maximum allowed size of " + maxSize + " bytes")); + return; + } + sink.next(buffer); + }); + }).collectList().map(buffers -> { + int totalSize = buffers.stream().mapToInt(ByteBuffer::remaining).sum(); + ByteBuffer combined = ByteBuffer.allocate(totalSize); + buffers.forEach(combined::put); + combined.flip(); + return StandardCharsets.UTF_8.decode(combined).toString(); + }).defaultIfEmpty(""); } - static class SseLineSubscriber extends BaseSubscriber { - - /** - * Pattern to extract data content from SSE "data:" lines. - */ - private static final Pattern EVENT_DATA_PATTERN = Pattern.compile("^data:(.+)$", Pattern.MULTILINE); - - /** - * Pattern to extract event ID from SSE "id:" lines. - */ - private static final Pattern EVENT_ID_PATTERN = Pattern.compile("^id:(.+)$", Pattern.MULTILINE); - - /** - * Pattern to extract event type from SSE "event:" lines. - */ - private static final Pattern EVENT_TYPE_PATTERN = Pattern.compile("^event:(.+)$", Pattern.MULTILINE); - - /** - * The sink for emitting parsed response events. - */ - private final FluxSink sink; - - /** - * StringBuilder for accumulating multi-line event data. - */ - private final StringBuilder eventBuilder; - - /** - * Current event's ID, if specified. - */ - private final AtomicReference currentEventId; - - /** - * Current event's type, if specified. - */ - private final AtomicReference currentEventType; - - /** - * The response information from the HTTP response. Send with each event to - * provide context. - */ - private ResponseInfo responseInfo; + /** + * Subscribes to the body publisher to release the underlying connection, discarding + * all bytes, then propagates the given error. + */ + static Flux drainThenError(Publisher> body, Throwable error) { + return JdkFlowAdapter.flowPublisherToFlux(body).thenMany(Flux.error(error)); + } - /** - * The maximum number of bytes that may accumulate for a single SSE event. A peer - * that never terminates an event (e.g. an endless stream of {@code data:} lines) - * has its stream aborted instead of exhausting memory. The accumulated data is - * measured in characters, which for UTF-8 is never more than the number of bytes - * it was decoded from. - */ - private final int maxSize; + /** + * Subscribes to the body publisher to release the underlying connection, discarding + * all bytes, then completes empty. + */ + static Flux drain(Publisher> body) { + return JdkFlowAdapter.flowPublisherToFlux(body).thenMany(Flux.empty()); + } - /** - * Creates a new LineSubscriber that will emit parsed SSE events to the provided - * sink. - * @param sink the {@link FluxSink} to emit parsed {@link ResponseEvent} objects - * to - * @param maxSize the maximum number of bytes that may accumulate for a single SSE - * event - */ - public SseLineSubscriber(ResponseInfo responseInfo, FluxSink sink, int maxSize) { - this.sink = sink; - this.eventBuilder = new StringBuilder(); - this.currentEventId = new AtomicReference<>(); - this.currentEventType = new AtomicReference<>(); - this.responseInfo = responseInfo; - this.maxSize = maxSize; - } + /** + * Stateful UTF-8 decoder that splits a stream of byte-buffer chunks into complete + * lines. Handles multi-byte characters split across chunk boundaries, and both + * {@code "\n"} and {@code "\r\n"} terminators. + */ + static final class Utf8LineDecoder { - @Override - protected void hookOnSubscribe(Subscription subscription) { + private final CharsetDecoder decoder = StandardCharsets.UTF_8.newDecoder(); - sink.onRequest(n -> { - subscription.request(n); - }); + private final CharBuffer charBuffer = CharBuffer.allocate(4096); - // Register disposal callback to cancel subscription when Flux is disposed - sink.onDispose(() -> { - subscription.cancel(); - }); - } + private final StringBuilder leftover = new StringBuilder(); - @Override - protected void hookOnNext(String line) { - if (line.isEmpty()) { - // Empty line means end of event - if (this.eventBuilder.length() > 0) { - String eventData = this.eventBuilder.toString(); - SseEvent sseEvent = new SseEvent(currentEventId.get(), currentEventType.get(), eventData.trim()); + // Holds partial UTF-8 sequences left over from a previous chunk (max 3 bytes + // for a BMP code point; 4 bytes for a supplementary one). + private ByteBuffer pendingBytes = ByteBuffer.allocate(0); - this.sink.next(new SseResponseEvent(responseInfo, sseEvent)); - this.eventBuilder.setLength(0); + List decode(List chunk) { + List lines = new ArrayList<>(); + for (ByteBuffer bb : chunk) { + ByteBuffer input = bb; + if (pendingBytes.hasRemaining()) { + ByteBuffer merged = ByteBuffer.allocate(pendingBytes.remaining() + bb.remaining()); + merged.put(pendingBytes).put(bb); + merged.flip(); + pendingBytes = ByteBuffer.allocate(0); + input = merged; } - } - else { - if (line.startsWith("data:")) { - var matcher = EVENT_DATA_PATTERN.matcher(line); - if (matcher.find()) { - String data = matcher.group(1).trim(); - // Measured before appending, so that an event carrying exactly - // maxSize of data is accepted: the trailing separator below is - // stripped again before the event is emitted. - if (this.eventBuilder.length() + data.length() > this.maxSize) { - upstream().cancel(); - this.sink.error( - new McpTransportException("Inbound SSE event exceeds the maximum allowed size of " - + this.maxSize + " bytes")); - return; + while (true) { + CoderResult result = decoder.decode(input, charBuffer, false); + drainCharBuffer(); + extractCompletedLines(lines); + if (result.isError()) { + try { + result.throwException(); + } + catch (CharacterCodingException e) { + throw new RuntimeException(e); } - this.eventBuilder.append(data).append("\n"); } - upstream().request(1); - } - else if (line.startsWith("id:")) { - var matcher = EVENT_ID_PATTERN.matcher(line); - if (matcher.find()) { - this.currentEventId.set(matcher.group(1).trim()); + if (result.isUnderflow()) { + if (input.hasRemaining()) { + pendingBytes = ByteBuffer.allocate(input.remaining()); + pendingBytes.put(input).flip(); + } + break; } - upstream().request(1); } - else if (line.startsWith("event:")) { - var matcher = EVENT_TYPE_PATTERN.matcher(line); - if (matcher.find()) { - this.currentEventType.set(matcher.group(1).trim()); - } - upstream().request(1); + } + return lines; + } + + List flush() { + ByteBuffer tail = pendingBytes.hasRemaining() ? pendingBytes : ByteBuffer.allocate(0); + CoderResult result = decoder.decode(tail, charBuffer, true); + while (result.isOverflow()) { + drainCharBuffer(); + result = decoder.decode(tail, charBuffer, true); + } + drainCharBuffer(); + if (result.isError()) { + try { + result.throwException(); } - else if (line.startsWith(":")) { - // Ignore comment lines starting with ":" - // This is a no-op, just to skip comments - logger.debug("Ignoring comment line: {}", line); - upstream().request(1); + catch (CharacterCodingException e) { + throw new RuntimeException(e); } - else { - // If the response is not successful, emit an error - this.sink.error(new McpTransportException( - "Invalid SSE response. Status code: " + this.responseInfo.statusCode() + " Line: " + line)); + } + result = decoder.flush(charBuffer); + while (result.isOverflow()) { + drainCharBuffer(); + result = decoder.flush(charBuffer); + } + drainCharBuffer(); + pendingBytes = ByteBuffer.allocate(0); + + List lines = new ArrayList<>(); + extractCompletedLines(lines); + if (leftover.length() > 0) { + String last = leftover.toString(); + if (last.endsWith("\r")) { + last = last.substring(0, last.length() - 1); } + leftover.setLength(0); + lines.add(last); } + return lines; } - @Override - protected void hookOnComplete() { - if (this.eventBuilder.length() > 0) { - String eventData = this.eventBuilder.toString(); - SseEvent sseEvent = new SseEvent(currentEventId.get(), currentEventType.get(), eventData.trim()); - this.sink.next(new SseResponseEvent(responseInfo, sseEvent)); - } - this.sink.complete(); + private void drainCharBuffer() { + charBuffer.flip(); + leftover.append(charBuffer); + charBuffer.clear(); } - @Override - protected void hookOnError(Throwable throwable) { - this.sink.error(throwable); + private void extractCompletedLines(List out) { + int newlineIdx; + while ((newlineIdx = leftover.indexOf("\n")) != -1) { + String line = leftover.substring(0, newlineIdx); + if (line.endsWith("\r")) { + line = line.substring(0, line.length() - 1); + } + out.add(line); + leftover.delete(0, newlineIdx + 1); + } } } - static class AggregateSubscriber extends BaseSubscriber { + /** + * Stateful SSE line parser. Accumulates {@code data:}, {@code id:} and {@code event:} + * fields until a blank line dispatches the event. Per the SSE spec, {@code id} and + * {@code event} persist across events until re-set; {@code data} is reset after each + * dispatch. + */ + static final class SseEventParser { - /** - * The sink for emitting parsed response events. - */ - private final FluxSink sink; + private static final Logger logger = LoggerFactory.getLogger(SseEventParser.class); - /** - * StringBuilder for accumulating multi-line event data. - */ - private final StringBuilder eventBuilder; + private final StringBuilder data = new StringBuilder(); /** - * The response information from the HTTP response. Send with each event to - * provide context. + * The maximum number of bytes that may accumulate for a single SSE event. A peer + * that never terminates an event (e.g. an endless stream of {@code data:} lines) + * has its stream aborted instead of exhausting memory. The accumulated data is + * measured in characters, which for UTF-8 is never more than the number of bytes + * it was decoded from. */ - private ResponseInfo responseInfo; + private final int maxSize; - volatile boolean hasRequestedDemand = false; + private String id; - /** - * The maximum number of bytes that may accumulate for the aggregated response - * body. A peer that sends a larger body has its response aborted instead of - * exhausting memory. The accumulated body is measured in characters, which for - * UTF-8 is never more than the number of bytes it was decoded from. - */ - private final int maxSize; + private String event; - /** - * Creates a new JsonLineSubscriber that will emit parsed JSON-RPC messages. - * @param sink the {@link FluxSink} to emit parsed {@link ResponseEvent} objects - * to - * @param maxSize the maximum number of bytes that may accumulate for the - * aggregated response body - */ - public AggregateSubscriber(ResponseInfo responseInfo, FluxSink sink, int maxSize) { - this.sink = sink; - this.eventBuilder = new StringBuilder(); - this.responseInfo = responseInfo; + SseEventParser(int maxSize) { this.maxSize = maxSize; } - @Override - protected void hookOnSubscribe(Subscription subscription) { - - sink.onRequest(n -> { - if (!hasRequestedDemand) { - subscription.request(Long.MAX_VALUE); + Optional feed(String line) { + if (line.isEmpty()) { + if (data.length() == 0) { + return Optional.empty(); } - hasRequestedDemand = true; - }); - - // Register disposal callback to cancel subscription when Flux is disposed - sink.onDispose(subscription::cancel); - } - - @Override - protected void hookOnNext(String line) { - // Measured before appending, so that a body of exactly maxSize is accepted. - // The separator this adds back for each line stands in for the terminator the - // peer sent, which the line subscriber has already stripped. - if (this.eventBuilder.length() + line.length() > this.maxSize) { - upstream().cancel(); - this.sink.error(new McpTransportException( - "Inbound response body exceeds the maximum allowed size of " + this.maxSize + " bytes")); - return; + SseEvent result = new SseEvent(id, event, data.toString().trim()); + data.setLength(0); + return Optional.of(result); } - this.eventBuilder.append(line).append("\n"); - } - - @Override - protected void hookOnComplete() { - - if (hasRequestedDemand) { - String data = this.eventBuilder.toString(); - this.sink.next(new AggregateResponseEvent(responseInfo, data)); + if (line.startsWith("data:")) { + String rest = line.substring(5); + if (!rest.isEmpty()) { + String value = rest.trim(); + // Measured before appending, so that an event carrying exactly + // maxSize of data is accepted: the trailing separator below is + // stripped again before the event is emitted. + if (data.length() + value.length() > this.maxSize) { + throw new McpTransportException( + "Inbound SSE event exceeds the maximum allowed size of " + this.maxSize + " bytes"); + } + data.append(value).append('\n'); + } } - - this.sink.complete(); - } - - @Override - protected void hookOnError(Throwable throwable) { - this.sink.error(throwable); - } - - } - - static class BodilessResponseLineSubscriber extends BaseSubscriber { - - /** - * The sink for emitting parsed response events. - */ - private final FluxSink sink; - - private final ResponseInfo responseInfo; - - volatile boolean hasRequestedDemand = false; - - public BodilessResponseLineSubscriber(ResponseInfo responseInfo, FluxSink sink) { - this.sink = sink; - this.responseInfo = responseInfo; - } - - @Override - protected void hookOnSubscribe(Subscription subscription) { - - sink.onRequest(n -> { - if (!hasRequestedDemand) { - subscription.request(Long.MAX_VALUE); + else if (line.startsWith("id:")) { + String rest = line.substring(3); + if (!rest.isEmpty()) { + id = rest.trim(); } - hasRequestedDemand = true; - }); - - // Register disposal callback to cancel subscription when Flux is disposed - sink.onDispose(() -> { - subscription.cancel(); - }); - } - - @Override - protected void hookOnComplete() { - if (hasRequestedDemand) { - // emit dummy event to be able to inspect the response info - // this is a shortcut allowing for a more streamlined processing using - // operator composition instead of having to deal with the - // CompletableFuture along the Subscriber for inspecting the result - this.sink.next(new DummyEvent(responseInfo)); } - this.sink.complete(); + else if (line.startsWith("event:")) { + String rest = line.substring(6); + if (!rest.isEmpty()) { + event = rest.trim(); + } + } + else if (line.startsWith(":")) { + logger.debug("Ignoring comment line: {}", line); + } + else { + throw new McpTransportException("Invalid SSE response line: " + line); + } + return Optional.empty(); } - @Override - protected void hookOnError(Throwable throwable) { - this.sink.error(throwable); + Optional flush() { + if (data.length() == 0) { + return Optional.empty(); + } + SseEvent result = new SseEvent(id, event, data.toString().trim()); + data.setLength(0); + return Optional.of(result); } } @@ -534,20 +463,24 @@ public void onComplete() { /** * A {@link BoundedBodySubscriber} that aborts the response once a single line (a run - * of bytes with no CR/LF terminator) exceeds {@code maxSize} bytes. + * of bytes with no LF) exceeds {@code maxSize} bytes. * *

- * {@link HttpResponse.BodySubscribers#fromLineSubscriber} buffers characters until it - * encounters a line terminator, so a peer that never terminates a line (or sends an - * enormous one) would force the transport to buffer it in memory. This wrapper counts - * bytes as they arrive off the wire and cancels the subscription before that buffer - * can grow without bound. + * {@link Utf8LineDecoder} buffers characters until it encounters a line terminator, + * so a peer that never terminates a line (or sends an enormous one) would force the + * transport to buffer it in memory. This wrapper counts bytes as they arrive off the + * wire and cancels the subscription before that buffer can grow without bound. + * + *

+ * Only LF resets the count, because LF is the only byte {@link Utf8LineDecoder} + * flushes a line on: a lone CR leaves the decoder's buffer growing, so it must not + * refill this budget either. CRLF still resets, on its LF. */ - static final class BoundedLineBodySubscriber extends BoundedBodySubscriber { + static final class BoundedLineBodySubscriber extends BoundedBodySubscriber { private long bytesSinceLineTerminator = 0; - BoundedLineBodySubscriber(BodySubscriber delegate, int maxSize) { + BoundedLineBodySubscriber(BodySubscriber delegate, int maxSize) { super(delegate, maxSize, "Inbound line"); } @@ -560,16 +493,16 @@ protected boolean checkSize(ByteBuffer buffer) { } if (this.bytesSinceLineTerminator + (limit - position) <= this.maxSize) { // No line ending in this buffer can exceed the limit, because there are - // not enough bytes since the last terminator for one to. Only the - // trailing (still unterminated) run matters, so scan back to the last - // terminator instead of walking every byte. + // not enough bytes since the last LF for one to. Only the trailing + // (still unterminated) run matters, so scan back to the last LF instead + // of walking every byte. this.bytesSinceLineTerminator = lengthOfTrailingRun(buffer, position, limit); return true; } // The limit is within reach, so account for every line exactly. for (int i = position; i < limit; i++) { byte b = buffer.get(i); - if (b == '\n' || b == '\r') { + if (b == '\n') { this.bytesSinceLineTerminator = 0; } else if (++this.bytesSinceLineTerminator > this.maxSize) { @@ -580,13 +513,12 @@ else if (++this.bytesSinceLineTerminator > this.maxSize) { } /** - * Returns the number of bytes after the last line terminator in the buffer, or - * the whole span added to the running count when the buffer holds no terminator. + * Returns the number of bytes after the last LF in the buffer, or the whole span + * added to the running count when the buffer holds no LF. */ private long lengthOfTrailingRun(ByteBuffer buffer, int position, int limit) { for (int i = limit - 1; i >= position; i--) { - byte b = buffer.get(i); - if (b == '\n' || b == '\r') { + if (buffer.get(i) == '\n') { return limit - 1 - i; } } diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java b/mcp-core/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java index fdb7bfd89..bfd71549f 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java @@ -78,6 +78,7 @@ public void close() { @Override public Mono closeGracefully() { return Mono.from(this.onClose.apply(this.sessionId.get())) + .onErrorResume(error -> Mono.fromRunnable(this.openConnections::dispose).then(Mono.error(error))) .then(Mono.fromRunnable(this.openConnections::dispose)); } diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/BoundedBodySubscriberTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/BoundedBodySubscriberTests.java index 5f350319e..af42b885f 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/BoundedBodySubscriberTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/BoundedBodySubscriberTests.java @@ -42,28 +42,28 @@ class BoundedBodySubscriberTests { @Test void lineSubscriberAcceptsEmptyBuffer() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); assertThat(subscriber.checkSize(buffer(""))).isTrue(); } @Test void lineSubscriberAcceptsLineOfExactlyMaxSize() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); assertThat(subscriber.checkSize(buffer("a".repeat(MAX_SIZE)))).isTrue(); } @Test void lineSubscriberRejectsLineOneByteOverMaxSize() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); assertThat(subscriber.checkSize(buffer("a".repeat(MAX_SIZE + 1)))).isFalse(); } @Test void lineSubscriberAccumulatesAcrossBuffers() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); assertThat(subscriber.checkSize(buffer("a".repeat(10)))).isTrue(); assertThat(subscriber.checkSize(buffer("a".repeat(6)))).isTrue(); @@ -73,7 +73,7 @@ void lineSubscriberAccumulatesAcrossBuffers() { @Test void lineSubscriberAcceptsUnboundedTotalOfTerminatedLines() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); // Far more than MAX_SIZE in total, but no single line comes close to it. for (int i = 0; i < 100; i++) { @@ -83,7 +83,7 @@ void lineSubscriberAcceptsUnboundedTotalOfTerminatedLines() { @Test void lineSubscriberResetsOnTerminatorAtEndOfBuffer() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); assertThat(subscriber.checkSize(buffer("a".repeat(10) + "\n"))).isTrue(); // A fresh line, so the previous 10 bytes must not count towards it. @@ -92,7 +92,7 @@ void lineSubscriberResetsOnTerminatorAtEndOfBuffer() { @Test void lineSubscriberResetsOnTerminatorAtStartOfBuffer() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); assertThat(subscriber.checkSize(buffer("a".repeat(MAX_SIZE)))).isTrue(); assertThat(subscriber.checkSize(buffer("\n" + "a".repeat(MAX_SIZE)))).isTrue(); @@ -100,15 +100,29 @@ void lineSubscriberResetsOnTerminatorAtStartOfBuffer() { @Test void lineSubscriberHandlesCrLfSplitAcrossBuffers() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); assertThat(subscriber.checkSize(buffer("a".repeat(12) + "\r"))).isTrue(); assertThat(subscriber.checkSize(buffer("\n" + "a".repeat(MAX_SIZE)))).isTrue(); } + @Test + void lineSubscriberIsNotResetByLoneCarriageReturns() { + BoundedLineBodySubscriber subscriber = lineSubscriber(); + + // Utf8LineDecoder only flushes a line on LF, so a peer streaming CR-terminated + // runs keeps its buffer growing. A lone CR must not refill the budget. + boolean accepted = true; + for (int i = 0; i < 10 && accepted; i++) { + accepted = subscriber.checkSize(buffer("aa\r")); + } + + assertThat(accepted).isFalse(); + } + @Test void lineSubscriberAcceptsBufferLargerThanMaxSizeHoldingOnlyShortLines() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); // Forces the exact per-byte accounting path: the buffer alone is well over the // limit, yet every line in it is legitimate. @@ -117,7 +131,7 @@ void lineSubscriberAcceptsBufferLargerThanMaxSizeHoldingOnlyShortLines() { @Test void lineSubscriberRejectsRunSpanningManyBuffers() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); boolean accepted = true; for (int i = 0; i < 10 && accepted; i++) { @@ -129,7 +143,7 @@ void lineSubscriberRejectsRunSpanningManyBuffers() { @Test void lineSubscriberOnlyCountsFromTheBufferPosition() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); ByteBuffer partiallyConsumed = buffer("a".repeat(MAX_SIZE * 2)); partiallyConsumed.position(MAX_SIZE * 2 - 4); @@ -138,7 +152,7 @@ void lineSubscriberOnlyCountsFromTheBufferPosition() { @Test void lineSubscriberDoesNotConsumeTheBuffer() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); ByteBuffer buffer = buffer("aaaa\nbbbb"); buffer.position(2); @@ -199,7 +213,7 @@ void totalSubscriberDoesNotConsumeTheBuffer() { @Test void forwardsBuffersWhileWithinBounds() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); List buffers = List.of(buffer("aaaa\n"), buffer("bbbb\n")); subscriber.onNext(buffers); @@ -211,7 +225,7 @@ void forwardsBuffersWhileWithinBounds() { @Test void abortsTheResponseWhenTheLineBoundIsExceeded() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); subscriber.onNext(List.of(buffer("a".repeat(MAX_SIZE + 1)))); @@ -235,7 +249,7 @@ void abortsTheResponseWhenTheTotalBoundIsExceeded() { @Test void withholdsTheWholeListWhenALaterBufferExceedsTheBound() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); subscriber.onNext(List.of(buffer("a".repeat(8)), buffer("a".repeat(9)))); @@ -245,7 +259,7 @@ void withholdsTheWholeListWhenALaterBufferExceedsTheBound() { @Test void ignoresFurtherSignalsOnceAborted() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); subscriber.onNext(List.of(buffer("a".repeat(MAX_SIZE + 1)))); Throwable firstError = this.delegate.error; @@ -262,8 +276,8 @@ void ignoresFurtherSignalsOnceAborted() { // --- fixtures ------------------------------------------------------------ - private BoundedLineBodySubscriber lineSubscriber() { - BoundedLineBodySubscriber subscriber = new BoundedLineBodySubscriber(this.delegate, MAX_SIZE); + private BoundedLineBodySubscriber lineSubscriber() { + BoundedLineBodySubscriber subscriber = new BoundedLineBodySubscriber<>(this.delegate, MAX_SIZE); subscriber.onSubscribe(this.subscription); return subscriber; } diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientLeakTestSupport.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientLeakTestSupport.java new file mode 100644 index 000000000..1bcce2466 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientLeakTestSupport.java @@ -0,0 +1,58 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.util.List; + +final class HttpClientLeakTestSupport { + + private HttpClientLeakTestSupport() { + } + + static int selectorManagerThreadCount() { + return selectorManagerThreadNames().size(); + } + + static List selectorManagerThreadNames() { + return Thread.getAllStackTraces() + .keySet() + .stream() + .map(Thread::getName) + .filter(name -> name.contains("HttpClient") && name.contains("SelectorManager")) + .sorted() + .toList(); + } + + static int forceGcUntilStable() throws InterruptedException { + int previousCount = Integer.MAX_VALUE; + int stableIterations = 0; + int currentCount = previousCount; + + for (int i = 0; i < 40; i++) { + System.gc(); + System.runFinalization(); + Thread.sleep(250); + + currentCount = selectorManagerThreadCount(); + if (currentCount == previousCount) { + stableIterations++; + if (stableIterations >= 4) { + break; + } + } + else { + stableIterations = 0; + previousCount = currentCount; + } + } + + return currentCount; + } + + static void pauseForSelectorStartup() throws InterruptedException { + Thread.sleep(150); + } + +} diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransportLeakTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransportLeakTests.java new file mode 100644 index 000000000..8b3ed2fd3 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransportLeakTests.java @@ -0,0 +1,51 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Function; + +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.json.gson.GsonMcpJsonMapper; +import org.junit.jupiter.api.Test; +import reactor.test.StepVerifier; + +import static org.assertj.core.api.Assertions.assertThat; + +class HttpClientSseClientTransportLeakTests { + + @Test + void closeDoesNotRetainOwnedHttpClientWhenClosedTransportRemainsReachable() throws Exception { + try (LoopbackMcpHttpServer server = LoopbackMcpHttpServer.start()) { + // List retainedClosedTransports = new + // ArrayList<>(); + int selectorThreadsBefore = HttpClientLeakTestSupport.selectorManagerThreadCount(); + Function, reactor.core.publisher.Mono> handler = Function + .identity(); + + for (int i = 0; i < 12; i++) { + HttpClientSseClientTransport transport = HttpClientSseClientTransport + .builder(server.baseUri().toString()) + .jsonMapper(new GsonMcpJsonMapper()) + .build(); + + StepVerifier.create(transport.connect(handler)).verifyComplete(); + HttpClientLeakTestSupport.pauseForSelectorStartup(); + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + // retainedClosedTransports.add(transport); + } + + int selectorThreadsAfter = HttpClientLeakTestSupport.forceGcUntilStable(); + + assertThat(selectorThreadsAfter) + .describedAs( + "closed transports should not keep owned HttpClient instances alive, remaining threads: %s", + HttpClientLeakTestSupport.selectorManagerThreadNames()) + .isLessThanOrEqualTo(selectorThreadsBefore + 1); + } + } + +} diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportLeakTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportLeakTests.java new file mode 100644 index 000000000..cdca6e9b6 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportLeakTests.java @@ -0,0 +1,55 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.util.Map; +import java.util.ArrayList; +import java.util.List; +import java.util.function.Function; + +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.json.gson.GsonMcpJsonMapper; +import org.junit.jupiter.api.Test; +import reactor.test.StepVerifier; + +import static org.assertj.core.api.Assertions.assertThat; + +class HttpClientStreamableHttpTransportLeakTests { + + @Test + void closeDoesNotRetainOwnedHttpClientWhenClosedTransportRemainsReachable() throws Exception { + try (LoopbackMcpHttpServer server = LoopbackMcpHttpServer.start()) { + // List retainedClosedTransports = new + // ArrayList<>(); + int selectorThreadsBefore = HttpClientLeakTestSupport.selectorManagerThreadCount(); + Function, reactor.core.publisher.Mono> handler = Function + .identity(); + + for (int i = 0; i < 12; i++) { + HttpClientStreamableHttpTransport transport = HttpClientStreamableHttpTransport + .builder(server.baseUri().toString()) + .jsonMapper(new GsonMcpJsonMapper()) + .build(); + + StepVerifier.create(transport.connect(handler)).verifyComplete(); + StepVerifier.create(transport.sendMessage( + new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, "ping", Map.of("iteration", i)))) + .verifyComplete(); + HttpClientLeakTestSupport.pauseForSelectorStartup(); + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + // retainedClosedTransports.add(transport); + } + + int selectorThreadsAfter = HttpClientLeakTestSupport.forceGcUntilStable(); + + assertThat(selectorThreadsAfter) + .describedAs( + "closed transports should not keep owned HttpClient instances alive, remaining threads: %s", + HttpClientLeakTestSupport.selectorManagerThreadNames()) + .isLessThanOrEqualTo(selectorThreadsBefore + 1); + } + } + +} diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/LoopbackMcpHttpServer.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/LoopbackMcpHttpServer.java new file mode 100644 index 000000000..01c7bb26d --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/LoopbackMcpHttpServer.java @@ -0,0 +1,136 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.net.URI; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadFactory; + +import com.sun.net.httpserver.Headers; +import com.sun.net.httpserver.HttpExchange; +import com.sun.net.httpserver.HttpHandler; +import com.sun.net.httpserver.HttpServer; + +final class LoopbackMcpHttpServer implements AutoCloseable { + + private static final byte[] EMPTY_BODY = new byte[0]; + + private static final byte[] STREAMABLE_PRIMER = """ + event: message + data: + + """.getBytes(StandardCharsets.UTF_8); + + private static final byte[] SSE_ENDPOINT_EVENT = """ + event: endpoint + data: /message + + """.getBytes(StandardCharsets.UTF_8); + + private final HttpServer server; + + private final ExecutorService executor; + + private LoopbackMcpHttpServer(HttpServer server, ExecutorService executor) { + this.server = server; + this.executor = executor; + } + + static LoopbackMcpHttpServer start() throws IOException { + HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0); + ExecutorService executor = Executors.newCachedThreadPool(new LoopbackThreadFactory()); + server.setExecutor(executor); + server.createContext("/mcp", new StreamableHandler()); + server.createContext("/sse", new SseHandler()); + server.createContext("/message", new MessageHandler()); + server.start(); + return new LoopbackMcpHttpServer(server, executor); + } + + URI baseUri() { + return URI.create("http://127.0.0.1:" + this.server.getAddress().getPort()); + } + + @Override + public void close() { + this.server.stop(0); + this.executor.shutdownNow(); + } + + private static final class StreamableHandler implements HttpHandler { + + @Override + public void handle(HttpExchange exchange) throws IOException { + try (exchange) { + String method = exchange.getRequestMethod(); + if ("GET".equals(method)) { + Headers headers = exchange.getResponseHeaders(); + headers.add("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, STREAMABLE_PRIMER.length); + try (OutputStream outputStream = exchange.getResponseBody()) { + outputStream.write(STREAMABLE_PRIMER); + } + return; + } + if ("POST".equals(method)) { + exchange.getResponseHeaders().add("mcp-session-id", "loopback-session"); + exchange.sendResponseHeaders(202, -1); + return; + } + if ("DELETE".equals(method)) { + exchange.sendResponseHeaders(204, -1); + return; + } + exchange.sendResponseHeaders(405, EMPTY_BODY.length); + } + } + + } + + private static final class SseHandler implements HttpHandler { + + @Override + public void handle(HttpExchange exchange) throws IOException { + try (exchange) { + Headers headers = exchange.getResponseHeaders(); + headers.add("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, SSE_ENDPOINT_EVENT.length); + try (OutputStream outputStream = exchange.getResponseBody()) { + outputStream.write(SSE_ENDPOINT_EVENT); + } + } + } + + } + + private static final class MessageHandler implements HttpHandler { + + @Override + public void handle(HttpExchange exchange) throws IOException { + try (exchange) { + exchange.sendResponseHeaders(202, -1); + } + } + + } + + private static final class LoopbackThreadFactory implements ThreadFactory { + + @Override + public Thread newThread(Runnable runnable) { + Thread thread = new Thread(runnable); + thread.setDaemon(true); + thread.setName("loopback-mcp-http-server-" + thread.getId()); + return thread; + } + + } + +} diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/SseEventParserTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/SseEventParserTests.java new file mode 100644 index 000000000..14272c11b --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/SseEventParserTests.java @@ -0,0 +1,120 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.util.Optional; + +import org.junit.jupiter.api.Test; + +import io.modelcontextprotocol.client.transport.ResponseSubscribers.SseEvent; +import io.modelcontextprotocol.client.transport.ResponseSubscribers.SseEventParser; +import io.modelcontextprotocol.spec.McpTransportException; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class SseEventParserTests { + + @Test + void simpleDataEvent() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed("data: hello")).isEmpty(); + Optional event = p.feed(""); + assertThat(event).isPresent(); + assertThat(event.get().data()).isEqualTo("hello"); + assertThat(event.get().id()).isNull(); + assertThat(event.get().event()).isNull(); + } + + @Test + void multiLineDataAccumulatesWithNewlineSeparatorAndTrims() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed("data: first")).isEmpty(); + assertThat(p.feed("data: second")).isEmpty(); + Optional event = p.feed(""); + assertThat(event).isPresent(); + assertThat(event.get().data()).isEqualTo("first\nsecond"); + } + + @Test + void idAndEventFieldsCaptured() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed("id: 42")).isEmpty(); + assertThat(p.feed("event: message")).isEmpty(); + assertThat(p.feed("data: payload")).isEmpty(); + Optional event = p.feed(""); + assertThat(event).isPresent(); + assertThat(event.get().id()).isEqualTo("42"); + assertThat(event.get().event()).isEqualTo("message"); + assertThat(event.get().data()).isEqualTo("payload"); + } + + @Test + void idAndEventPersistAcrossEvents() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + p.feed("id: 1"); + p.feed("event: message"); + p.feed("data: one"); + SseEvent first = p.feed("").orElseThrow(); + assertThat(first.id()).isEqualTo("1"); + assertThat(first.event()).isEqualTo("message"); + + p.feed("data: two"); + SseEvent second = p.feed("").orElseThrow(); + assertThat(second.id()).isEqualTo("1"); + assertThat(second.event()).isEqualTo("message"); + assertThat(second.data()).isEqualTo("two"); + } + + @Test + void commentLineIgnored() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed(": this is a comment")).isEmpty(); + assertThat(p.feed("data: hello")).isEmpty(); + Optional event = p.feed(""); + assertThat(event).isPresent(); + assertThat(event.get().data()).isEqualTo("hello"); + } + + @Test + void trailingIncompleteEventEmittedOnFlush() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + p.feed("data: incomplete"); + Optional flushed = p.flush(); + assertThat(flushed).isPresent(); + assertThat(flushed.get().data()).isEqualTo("incomplete"); + } + + @Test + void flushWithNothingPendingIsEmpty() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.flush()).isEmpty(); + } + + @Test + void blankLineWithNoPendingDataIsEmpty() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed("")).isEmpty(); + } + + @Test + void unknownFieldThrowsMcpTransportException() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThatThrownBy(() -> p.feed("bogus line")).isInstanceOf(McpTransportException.class) + .hasMessageContaining("Invalid SSE response line") + .hasMessageContaining("bogus line"); + } + + @Test + void dataFieldWithEmptyValueIsIgnored() { + // preserves the pre-refactor regex behavior where `^data:(.+)$` required at + // least one char after the colon — a lone `data:` line contributes nothing. + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed("data:")).isEmpty(); + // no pending data → blank line dispatches nothing + assertThat(p.feed("")).isEmpty(); + } + +} diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java new file mode 100644 index 000000000..c321fc226 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java @@ -0,0 +1,135 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.List; + +import org.junit.jupiter.api.Test; + +import io.modelcontextprotocol.client.transport.ResponseSubscribers.Utf8LineDecoder; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +class Utf8LineDecoderTests { + + private static List chunk(String... parts) { + return List.of(toByteBuffers(parts)); + } + + private static ByteBuffer[] toByteBuffers(String... parts) { + ByteBuffer[] bbs = new ByteBuffer[parts.length]; + for (int i = 0; i < parts.length; i++) { + bbs[i] = ByteBuffer.wrap(parts[i].getBytes(StandardCharsets.UTF_8)); + } + return bbs; + } + + @Test + void singleLineLf() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("hello\n"))).containsExactly("hello"); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void singleLineCrLf() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("hello\r\n"))).containsExactly("hello"); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void multipleLinesInOneChunk() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("one\ntwo\nthree\n"))).containsExactly("one", "two", "three"); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void lineSplitAcrossChunks() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("hel"))).isEmpty(); + assertThat(dec.decode(chunk("lo\nworld"))).containsExactly("hello"); + assertThat(dec.flush()).containsExactly("world"); + } + + @Test + void lineSplitAcrossByteBuffersInSameChunk() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + // two byte-buffers, newline between them -- should still form one clean split + List chunk = List.of(ByteBuffer.wrap("part-one\npart-".getBytes(StandardCharsets.UTF_8)), + ByteBuffer.wrap("two\n".getBytes(StandardCharsets.UTF_8))); + assertThat(new Utf8LineDecoder().decode(chunk)).containsExactly("part-one", "part-two"); + } + + @Test + void multiByteUtf8SplitAcrossChunks() { + // "€" is U+20AC → 0xE2 0x82 0xAC in UTF-8. Split between the first and second + // byte. + byte[] euro = "€".getBytes(StandardCharsets.UTF_8); + assertThat(euro).hasSize(3); + + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(List.of(ByteBuffer.wrap(new byte[] { euro[0] })))).isEmpty(); + assertThat(dec.decode(List.of(ByteBuffer.wrap(new byte[] { euro[1], euro[2], '\n' })))).containsExactly("€"); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void consecutiveBlankLines() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("\n\n\n"))).containsExactly("", "", ""); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void trailingPartialLineEmittedOnFlush() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("incomplete"))).isEmpty(); + assertThat(dec.flush()).containsExactly("incomplete"); + } + + @Test + void trailingPartialLineWithCrTrimmedOnFlush() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("incomplete\r"))).isEmpty(); + assertThat(dec.flush()).containsExactly("incomplete"); + } + + @Test + void emptyInput() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(List.of())).isEmpty(); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void invalidUtf8Throws() { + // 0xFF is never a valid UTF-8 lead byte + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThatThrownBy(() -> dec.decode(List.of(ByteBuffer.wrap(new byte[] { (byte) 0xFF })))) + .isInstanceOf(RuntimeException.class); + } + + @Test + void linesLongerThanInternalCharBuffer() { + // 4096 is the internal CharBuffer size; send a single line ~10k chars to force + // multiple overflow cycles. + StringBuilder big = new StringBuilder(); + for (int i = 0; i < 10_000; i++) { + big.append('a'); + } + big.append('\n'); + + Utf8LineDecoder dec = new Utf8LineDecoder(); + List lines = dec.decode(chunk(big.toString())); + assertThat(lines).hasSize(1); + assertThat(lines.get(0)).hasSize(10_000); + } + +} diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/spec/DefaultMcpTransportSessionTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/spec/DefaultMcpTransportSessionTests.java new file mode 100644 index 000000000..a8ffac41e --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/spec/DefaultMcpTransportSessionTests.java @@ -0,0 +1,26 @@ +package io.modelcontextprotocol.spec; + +import java.util.concurrent.atomic.AtomicBoolean; + +import org.junit.jupiter.api.Test; +import reactor.core.Disposable; +import reactor.core.publisher.Mono; +import reactor.test.StepVerifier; + +import static org.assertj.core.api.Assertions.assertThat; + +class DefaultMcpTransportSessionTests { + + @Test + void closeGracefullyDisposesOpenConnectionsEvenWhenOnCloseFails() { + var disposed = new AtomicBoolean(); + Disposable disposable = () -> disposed.set(true); + var session = new DefaultMcpTransportSession(id -> Mono.error(new RuntimeException("boom"))); + session.addConnection(disposable); + + StepVerifier.create(session.closeGracefully()).expectErrorMessage("boom").verify(); + + assertThat(disposed.get()).isTrue(); + } + +} diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientBoundedReadTestSupport.java b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientBoundedReadTestSupport.java index a22dc1301..b962b55ab 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientBoundedReadTestSupport.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientBoundedReadTestSupport.java @@ -97,6 +97,23 @@ protected static Responder unterminatedLine(int chunks) { }; } + /** + * 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. + */ + 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(); + } + }; + } + /** * A responder that writes enough short, properly terminated lines to exceed the limit * in aggregate. diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportBoundedReadTests.java b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportBoundedReadTests.java index 856dc88fb..bc47f8540 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportBoundedReadTests.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportBoundedReadTests.java @@ -40,6 +40,16 @@ void shouldRejectSingleLineExceedingMaxSize() { .verifyErrorMatches(t -> messageContains(t, "Inbound line exceeds the maximum allowed size")); } + @Test + void shouldRejectCarriageReturnTerminatedRunsExceedingMaxSize() { + // The line decoder only flushes on LF, so CR-terminated runs keep its buffer + // growing. A lone CR must not be treated as the start of a fresh line's budget. + respondWith(endpoint(), "text/event-stream", carriageReturnTerminatedRuns(8)); + + StepVerifier.create(sendMessage()) + .verifyErrorMatches(t -> messageContains(t, "Inbound line exceeds the maximum allowed size")); + } + @Test void shouldRejectEventExceedingMaxSize() { // Many short, terminated "data:" lines with no blank line to end the event. Each From f90bcd6951c0f0f180eb664ce20f8aad5031e60d Mon Sep 17 00:00:00 2001 From: Daniel Garnier-Moiroux Date: Thu, 27 Aug 2026 16:46:13 +0200 Subject: [PATCH 2/2] Improve ResponseSubscribers performance Fixes #1042 Signed-off-by: Daniel Garnier-Moiroux --- .../client/transport/ResponseSubscribers.java | 18 +- .../transport/LargeSseEventDecodingTests.java | 219 +++++++++++++ .../transport/Utf8LineDecoderTests.java | 30 ++ ...amableHttpTransportLargeResponseTests.java | 295 ++++++++++++++++++ 4 files changed, 561 insertions(+), 1 deletion(-) create mode 100644 mcp-core/src/test/java/io/modelcontextprotocol/client/transport/LargeSseEventDecodingTests.java create mode 100644 mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportLargeResponseTests.java diff --git a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java index 202f525ca..e492c72ff 100644 --- a/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java +++ b/mcp-core/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java @@ -191,6 +191,17 @@ static final class Utf8LineDecoder { private final StringBuilder leftover = new StringBuilder(); + /** + * How many leading characters of {@link #leftover} are already known to hold no + * line terminator, so that the search for one resumes where the previous search + * ended instead of restarting at the beginning of the buffer. Without it, a long + * line is searched again in full for every chunk that arrives, which makes + * reading an event cost time proportional to the square of its length. + * @see #1042 + */ + private int scannedForLineTerminator = 0; + // Holds partial UTF-8 sequences left over from a previous chunk (max 3 bytes // for a BMP code point; 4 bytes for a supplementary one). private ByteBuffer pendingBytes = ByteBuffer.allocate(0); @@ -263,6 +274,7 @@ List flush() { last = last.substring(0, last.length() - 1); } leftover.setLength(0); + this.scannedForLineTerminator = 0; lines.add(last); } return lines; @@ -276,14 +288,18 @@ private void drainCharBuffer() { private void extractCompletedLines(List out) { int newlineIdx; - while ((newlineIdx = leftover.indexOf("\n")) != -1) { + while ((newlineIdx = leftover.indexOf("\n", this.scannedForLineTerminator)) != -1) { String line = leftover.substring(0, newlineIdx); if (line.endsWith("\r")) { line = line.substring(0, line.length() - 1); } out.add(line); leftover.delete(0, newlineIdx + 1); + // What is left starts after the terminator, so none of it has been + // searched yet. + this.scannedForLineTerminator = 0; } + this.scannedForLineTerminator = leftover.length(); } } diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/LargeSseEventDecodingTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/LargeSseEventDecodingTests.java new file mode 100644 index 000000000..bd4b641fc --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/LargeSseEventDecodingTests.java @@ -0,0 +1,219 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.Flow; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.adapter.JdkFlowAdapter; +import reactor.core.publisher.Flux; + +import io.modelcontextprotocol.client.transport.ResponseSubscribers.SseEvent; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * Reproducer for the client-side SSE reading bottleneck reported in + * #1042: a + * tool response arriving as a single multi-megabyte {@code data:} line, which is what + * compact JSON looks like on the wire, took ~5s to read where the same bytes read with + * {@link java.net.http.HttpResponse.BodyHandlers#ofString()} took ~0.4s. + * + *

+ * What cost the time was the length of the line rather than the number of bytes, because + * the buffered characters were gone over again every time a chunk arrived. + * {@link #shouldReadOneLargeEventWithinBudgetOfManySmallOnes()} therefore measures the + * same payload twice, once as one long line and once split over short ones, and asserts a + * ratio rather than a duration, so that it keeps its meaning on a machine of any speed. + * + *

+ * Reading a single-line event in 16KiB chunks, best of three runs on the same machine: + * + * 

+ * payload   2.0.0 (fromLineSubscriber)   rescanning the line   scanning each line once
+ *  1MiB                        210ms                    12ms                      2ms
+ *  2MiB                        810ms                    32ms                      6ms
+ *  4MiB                       3314ms                   138ms                     10ms
+ *  8MiB                      13140ms                   561ms                     18ms
+ *
+ * + *

+ * 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 + * 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 + * brings the ratio this test measures from ~12 to ~1.6. + * + *

+ * See {@code HttpClientStreamableHttpTransportLargeResponseTests} in {@code mcp-test} for + * the same comparison end to end, over a real connection. + * + * @author Daniel Garnier-Moiroux + */ +class LargeSseEventDecodingTests { + + private static final Logger logger = LoggerFactory.getLogger(LargeSseEventDecodingTests.class); + + /** + * Roughly what {@link java.net.http.HttpClient} hands to a body subscriber at a time. + * The cost the report is about was paid per chunk, so the chunking is part of the + * reproducer. + */ + private static final int CHUNK_SIZE = 16 * 1024; + + private static final int MIB = 1024 * 1024; + + /** + * The payload size in the report: ~4MiB of compact JSON, and therefore ~4MiB with no + * line terminator in it. + */ + private static final int PAYLOAD_SIZE = 4 * MIB; + + /** + * How much of {@link #PAYLOAD_SIZE} each event carries when the same total is split + * over many events. + */ + private static final int SMALL_EVENT_SIZE = 64 * 1024; + + /** + * How much longer decoding the payload as one long line may take than decoding the + * same bytes as short ones. A reader that goes over each line once is indifferent to + * how long the lines are, which measures ~1.6 here; the bound leaves headroom over + * that, and is far below what rescanning the line measured (~12x) or what 2.0.0 + * measured (~220x). + */ + private static final double MAX_SINGLE_EVENT_PENALTY = 4.0; + + private static final int MAX_SIZE = 64 * MIB; + + @Test + @Timeout(60) + void shouldDecodeMultiMegabyteSingleLineEventIntact() { + String payload = payloadOfSize(PAYLOAD_SIZE); + + List events = decode(oneLargeEvent(payload)); + + assertThat(events).hasSize(1); + assertThat(events.get(0).event()).isEqualTo("message"); + assertThat(events.get(0).data()).isEqualTo(payload); + } + + @Test + @Timeout(300) + void shouldReadOneLargeEventWithinBudgetOfManySmallOnes() { + byte[] oneEvent = oneLargeEvent(payloadOfSize(PAYLOAD_SIZE)); + byte[] manyEvents = manySmallEvents(PAYLOAD_SIZE, SMALL_EVENT_SIZE); + int smallEventCount = PAYLOAD_SIZE / SMALL_EVENT_SIZE; + + // The reporter measured a JIT effect at this payload size: the first few large + // reads spike before the hot loop settles. Warm up, then interleave the two + // shapes and take the best of each, so the comparison reflects steady state. + for (int i = 0; i < 5; i++) { + decode(manyEvents); + } + long single = Long.MAX_VALUE; + long split = Long.MAX_VALUE; + for (int i = 0; i < 3; i++) { + single = Math.min(single, timeDecode(oneEvent, 1)); + split = Math.min(split, timeDecode(manyEvents, smallEventCount)); + } + + double penalty = (double) single / Math.max(split, 1); + logger.info("decoded {}KiB as one event in {}ms and as {} events in {}ms: ratio {}", PAYLOAD_SIZE / 1024, + single / 1_000_000, smallEventCount, split / 1_000_000, String.format("%.1f", penalty)); + logScaling(); + + assertThat(penalty) + .as("decoding %dKiB as a single SSE event took %.1fx as long as decoding the same number of bytes as " + + "%dKiB events, so the cost of an event grows with the length of its line", PAYLOAD_SIZE / 1024, + penalty, SMALL_EVENT_SIZE / 1024) + .isLessThan(MAX_SINGLE_EVENT_PENALTY); + } + + /** + * Logs how decoding one long line scales with its length, which is the shape the + * report is about: doubling the payload should cost about twice the time, not four + * times it. Not asserted, because the ratio above covers the same ground with a + * baseline measured on the same machine. + */ + private void logScaling() { + for (int payloadSize : new int[] { MIB, 2 * MIB, 4 * MIB, 8 * MIB }) { + byte[] body = oneLargeEvent(payloadOfSize(payloadSize)); + long best = Math.min(timeDecode(body, 1), timeDecode(body, 1)); + logger.info("decoded a single-line event of {}KiB in {}ms", payloadSize / 1024, best / 1_000_000); + } + } + + /** + * Decodes the body once and returns how long it took, in nanoseconds. + */ + private static long timeDecode(byte[] body, int expectedEvents) { + long start = System.nanoTime(); + List events = decode(body); + long elapsed = System.nanoTime() - start; + assertThat(events).hasSize(expectedEvents); + return elapsed; + } + + /** + * Runs the body through the transport's SSE reading path, as + * {@code HttpClientStreamableHttpTransport} does, chunked the way the HTTP client + * chunks a response body. + */ + private static List decode(byte[] body) { + Flow.Publisher> publisher = JdkFlowAdapter + .publisherToFlowPublisher(Flux.fromIterable(chunk(body))); + Flux lines = ResponseSubscribers.decodeLines(publisher); + return ResponseSubscribers.decodeSseResponse(lines, MAX_SIZE).collectList().block(); + } + + private static List> chunk(byte[] body) { + List> chunks = new ArrayList<>(); + for (int offset = 0; offset < body.length; offset += CHUNK_SIZE) { + int length = Math.min(CHUNK_SIZE, body.length - offset); + chunks.add(List.of(ByteBuffer.wrap(body, offset, length).asReadOnlyBuffer())); + } + return chunks; + } + + /** + * An SSE {@code message} event carrying the whole payload on a single {@code data:} + * line. + */ + private static byte[] oneLargeEvent(String payload) { + return ("event: message\ndata: " + payload + "\n\n").getBytes(StandardCharsets.UTF_8); + } + + /** + * The same {@code total} number of payload bytes, spread over events of + * {@code eachSize} each. + */ + private static byte[] manySmallEvents(int total, int eachSize) { + StringBuilder body = new StringBuilder(total + 4096); + for (int i = 0; i < total / eachSize; i++) { + body.append("event: message\ndata: ").append(payloadOfSize(eachSize)).append("\n\n"); + } + return body.toString().getBytes(StandardCharsets.UTF_8); + } + + /** + * A single-line JSON-RPC response of exactly {@code size} characters, none of them a + * line terminator. + */ + private static String payloadOfSize(int size) { + String prefix = "{\"jsonrpc\":\"2.0\",\"id\":\"test-id\",\"result\":{\"content\":\""; + String suffix = "\"}}"; + return prefix + "a".repeat(size - prefix.length() - suffix.length()) + suffix; + } + +} diff --git a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java index c321fc226..51838b95a 100644 --- a/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java @@ -116,6 +116,36 @@ void invalidUtf8Throws() { .isInstanceOf(RuntimeException.class); } + @Test + void resumesSearchAfterTerminatorWhenLineWasSplitAcrossManyChunks() { + // The decoder remembers how far it has searched for a terminator, so the chunk + // that finally terminates a long line must not leave that mark behind and hide + // the lines that follow it. + Utf8LineDecoder dec = new Utf8LineDecoder(); + for (int i = 0; i < 10; i++) { + assertThat(dec.decode(chunk("aaaa"))).isEmpty(); + } + assertThat(dec.decode(chunk("\nsecond\nthird\n"))).containsExactly("a".repeat(40), "second", "third"); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void resumesSearchAcrossChunkWhenTerminatorFollowsUnterminatedPrefix() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("unterminated"))).isEmpty(); + assertThat(dec.decode(chunk("-still-going"))).isEmpty(); + assertThat(dec.decode(chunk("\n"))).containsExactly("unterminated-still-going"); + } + + @Test + void crLfSplitAcrossChunks() { + // The CR ends one chunk and the LF opens the next, so the terminator sits exactly + // at the point the previous search stopped. + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("hello\r"))).isEmpty(); + assertThat(dec.decode(chunk("\nworld\r\n"))).containsExactly("hello", "world"); + } + @Test void linesLongerThanInternalCharBuffer() { // 4096 is the internal CharBuffer size; send a single line ~10k chars to force diff --git a/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportLargeResponseTests.java b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportLargeResponseTests.java new file mode 100644 index 000000000..51c62997b --- /dev/null +++ b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportLargeResponseTests.java @@ -0,0 +1,295 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.io.IOException; +import java.io.OutputStream; +import java.net.InetSocketAddress; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Consumer; + +import com.sun.net.httpserver.HttpServer; +import io.modelcontextprotocol.server.transport.TomcatTestUtil; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage; +import io.modelcontextprotocol.spec.McpSchema.JSONRPCNotification; +import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest; +import io.modelcontextprotocol.spec.McpSchema.JSONRPCResponse; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import reactor.core.publisher.Mono; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * End-to-end reproducer for + * #1042: a + * tool result of a few megabytes, delivered on the POST response's SSE stream as a single + * compact-JSON {@code data:} line, took the client ~5s to read where {@code curl} and + * {@link java.net.http.HttpResponse.BodyHandlers#ofString()} read the same bytes in + * ~0.4s. + * + *

+ * The reported setup is reproduced with a bare {@link HttpServer} in place of an MCP + * server, so that only the client's reading of the response is measured. What makes the + * payload expensive is that it arrives as one very long line, because compact JSON has no + * newline in it, so + * {@link #shouldReadOneLargeEventAboutAsFastAsTheSameBytesSplitOverManyEvents()} compares + * reading it against reading the same number of bytes split over many short events. That + * ratio is what the line length costs, with everything else (the wire, the JSON parsing, + * the machine) held constant. + * + *

+ * Measured here for a 4MiB payload, best of 12 reads each: ~26ms as one event against + * ~26ms split up, a ratio of ~1. Before the line decoder resumed its search for a line + * terminator where the previous search ended, instead of restarting it for every chunk + * that arrives, the same comparison measured ~180ms against ~30ms, and 2.0.0 measured + * ~100x. See {@code LargeSseEventDecodingTests} in {@code mcp-core} for the same + * comparison without a wire in between, and for how it scales with the payload. + * + *

+ * The per-read timings are logged. The issue also reported the first few reads of a large + * response taking an order of magnitude longer than the ones after them, which was that + * per-chunk cost paid while the hot loop was still being compiled: at this payload size + * the reads now go 404ms, 48ms, 39ms, then settle at ~26ms. The assertion is still made + * on the best read of many, so that it describes steady state rather than compilation. + * + * @author Daniel Garnier-Moiroux + */ +@Timeout(300) +class HttpClientStreamableHttpTransportLargeResponseTests { + + private static final Logger logger = LoggerFactory + .getLogger(HttpClientStreamableHttpTransportLargeResponseTests.class); + + private static final String ENDPOINT = "/mcp"; + + private static final String REQUEST_ID = "test-id"; + + /** + * The payload size in the report: ~4MiB of compact JSON, and therefore ~4MiB with no + * line terminator in it. + */ + private static final int PAYLOAD_SIZE = 4 * 1024 * 1024; + + /** + * How much of {@link #PAYLOAD_SIZE} each event carries when the same total is split + * over many events. + */ + private static final int SMALL_EVENT_SIZE = 64 * 1024; + + private static final int MEASURED_READS = 12; + + /** + * How much longer reading the payload as one event may take than reading the same + * bytes split over many events. A reader that goes over each line once is indifferent + * to how long the lines are, which measures ~1 here; the bound leaves headroom over + * that for a loaded machine, and is far below what rescanning the line measured + * (~6x). + */ + private static final double MAX_SINGLE_EVENT_PENALTY = 2.5; + + /** + * Writes the SSE body of the POST response, and may block until the client has + * reacted to what it has written so far. + */ + @FunctionalInterface + private interface SseResponder { + + void respond(OutputStream body) throws IOException, InterruptedException; + + } + + private HttpServer server; + + private String host; + + private volatile SseResponder responder; + + @BeforeEach + void startServer() throws IOException { + int port = TomcatTestUtil.findAvailablePort(); + this.host = "http://localhost:" + port; + this.server = HttpServer.create(new InetSocketAddress(port), 0); + this.server.setExecutor(Executors.newCachedThreadPool()); + this.server.createContext(ENDPOINT, exchange -> { + try (exchange) { + if (!"POST".equals(exchange.getRequestMethod())) { + // The transport opens a server-initiated stream after its first POST. + // 405 tells it there is none, which keeps this fixture to a single + // request-response exchange. + exchange.sendResponseHeaders(405, -1); + return; + } + exchange.getRequestBody().readAllBytes(); + exchange.getResponseHeaders().set("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, 0); + try (OutputStream body = exchange.getResponseBody()) { + this.responder.respond(body); + } + catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } + } + }); + this.server.start(); + } + + @AfterEach + void stopServer() { + if (this.server != null) { + this.server.stop(0); + } + } + + @Test + void shouldReceiveLargeSingleLineSseResponseIntact() throws Exception { + String payload = "a".repeat(PAYLOAD_SIZE); + this.responder = body -> writeEvent(body, jsonRpcResponse(payload)); + + List received = new CopyOnWriteArrayList<>(); + long elapsed = time(() -> readResponse(received::add)); + logger.info("received a {}KiB single-line SSE response in {}ms", PAYLOAD_SIZE / 1024, elapsed / 1_000_000); + + assertThat(received).hasSize(1); + JSONRPCResponse response = (JSONRPCResponse) received.get(0); + assertThat(response.id()).isEqualTo(REQUEST_ID); + assertThat(((Map) response.result()).get("content")).isEqualTo(payload); + } + + @Test + void shouldDeliverInterleavedNotificationBeforeTheLargeResultIsWritten() throws Exception { + // The response interleaves a progress notification before the result, on the same + // stream, which is why the issue rules out reading the whole body in one go: each + // event has to be delivered as its boundary arrives. This server refuses to write + // the result until the client has acknowledged the notification, so a reader that + // waits for the whole body deadlocks instead of quietly passing. + CountDownLatch notificationDelivered = new CountDownLatch(1); + AtomicBoolean deliveredBeforeResult = new AtomicBoolean(); + this.responder = body -> { + writeEvent(body, progressNotification("")); + deliveredBeforeResult.set(notificationDelivered.await(60, TimeUnit.SECONDS)); + writeEvent(body, jsonRpcResponse("a".repeat(PAYLOAD_SIZE))); + }; + + List received = new CopyOnWriteArrayList<>(); + readResponse(message -> { + received.add(message); + if (message instanceof JSONRPCNotification) { + notificationDelivered.countDown(); + } + }); + + assertThat(deliveredBeforeResult) + .as("the notification sent before the %dKiB result was not delivered until the whole response body had been read", + PAYLOAD_SIZE / 1024) + .isTrue(); + assertThat(received).hasSize(2); + assertThat(received.get(0)).isInstanceOf(JSONRPCNotification.class); + assertThat(received.get(1)).isInstanceOf(JSONRPCResponse.class); + } + + @Test + void shouldReadOneLargeEventWithinBudgetOfManySmallOnes() { + String payload = "a".repeat(PAYLOAD_SIZE); + String chunk = "a".repeat(SMALL_EVENT_SIZE); + SseResponder oneLargeEvent = body -> writeEvent(body, jsonRpcResponse(payload)); + SseResponder manySmallEvents = body -> { + for (int i = 0; i < PAYLOAD_SIZE / SMALL_EVENT_SIZE; i++) { + writeEvent(body, progressNotification(chunk)); + } + writeEvent(body, jsonRpcResponse("")); + }; + + // Interleaved, so that both shapes see the same machine and the same JIT state. + long oneEvent = Long.MAX_VALUE; + long manyEvents = Long.MAX_VALUE; + for (int i = 0; i < MEASURED_READS; i++) { + this.responder = oneLargeEvent; + long oneEventRead = time(() -> readResponse(message -> { + })); + this.responder = manySmallEvents; + long manyEventsRead = time(() -> readResponse(message -> { + })); + logger.info("read #{} of {}KiB: {}ms as one event, {}ms split over {}KiB events", i + 1, + PAYLOAD_SIZE / 1024, oneEventRead / 1_000_000, manyEventsRead / 1_000_000, SMALL_EVENT_SIZE / 1024); + oneEvent = Math.min(oneEvent, oneEventRead); + manyEvents = Math.min(manyEvents, manyEventsRead); + } + + double penalty = (double) oneEvent / Math.max(manyEvents, 1); + logger.info("best read: {}ms as one event, {}ms split up: ratio {}", oneEvent / 1_000_000, + manyEvents / 1_000_000, String.format("%.1f", penalty)); + + assertThat(penalty) + .as("reading %dKiB as a single SSE event took %.1fx as long as reading the same number of bytes split " + + "over %dKiB events, so the cost of an event grows with the length of its line", + PAYLOAD_SIZE / 1024, penalty, SMALL_EVENT_SIZE / 1024) + .isLessThan(MAX_SINGLE_EVENT_PENALTY); + } + + /** + * Sends one request and returns once the response has been delivered, handing every + * message received on the way to {@code onMessage}. + */ + private void readResponse(Consumer onMessage) { + HttpClientStreamableHttpTransport transport = HttpClientStreamableHttpTransport.builder(this.host) + .endpoint(ENDPOINT) + .build(); + CompletableFuture response = new CompletableFuture<>(); + JSONRPCRequest request = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "tools/call", REQUEST_ID, + Map.of("name", "large-response")); + try { + transport.connect(messages -> messages.doOnNext(message -> { + onMessage.accept(message); + if (message instanceof JSONRPCResponse) { + response.complete(message); + } + })).then(transport.sendMessage(request)).block(Duration.ofSeconds(120)); + response.get(120, TimeUnit.SECONDS); + } + catch (Exception e) { + throw new RuntimeException(e); + } + finally { + transport.closeGracefully().block(Duration.ofSeconds(10)); + } + } + + private static long time(Runnable read) { + long start = System.nanoTime(); + read.run(); + return System.nanoTime() - start; + } + + private static void writeEvent(OutputStream body, String data) throws IOException { + body.write(("event: message\ndata: " + data + "\n\n").getBytes(StandardCharsets.UTF_8)); + body.flush(); + } + + private static String jsonRpcResponse(String payload) { + return "{\"jsonrpc\":\"2.0\",\"id\":\"" + REQUEST_ID + "\",\"result\":{\"content\":\"" + payload + "\"}}"; + } + + private static String progressNotification(String payload) { + return "{\"jsonrpc\":\"2.0\",\"method\":\"notifications/progress\",\"params\":{\"progressToken\":\"" + + REQUEST_ID + "\",\"progress\":1,\"total\":2,\"message\":\"" + payload + "\"}}"; + } + +}