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..faaf59beb 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,68 @@ 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(); + // The body is handed over as a publisher and nothing is read off + // the wire until it is subscribed, so it has to be drained even + // when its content is of no further interest. + return ResponseSubscribers.drain(response.body()); } - 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.error(new RuntimeException("Failed to handle SSE endpoint event")); + } + } + else if (MESSAGE_EVENT_TYPE.equals(sseEvent.event())) { + String data = sseEvent.data(); + if (data == null || data.isBlank()) { + logger.debug("Skipping SSE event with empty data (stream primer)"); sink.success(); + return Flux.empty(); } + JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(jsonMapper, data); + sink.success(); + return Flux.just(message); } - catch (IOException e) { - sink.error(new McpTransportException("Error processing SSE event", e)); + 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 +534,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..c530f8d5f 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,44 @@ 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); + if (onFirstMessage != null) { + onFirstMessage.run(); + } + 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 +342,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 +367,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 +455,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 +483,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 +512,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 +520,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 +533,117 @@ 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(); + 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))) { + 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..907d77eb2 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,41 @@ /* -* 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.CodingErrorAction; +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 +43,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 +55,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 +104,335 @@ 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); + /** + * 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)); + } - /** - * Pattern to extract event ID from SSE "id:" lines. - */ - private static final Pattern EVENT_ID_PATTERN = Pattern.compile("^id:(.+)$", Pattern.MULTILINE); + /** + * 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()); + } - /** - * Pattern to extract event type from SSE "event:" lines. - */ - private static final Pattern EVENT_TYPE_PATTERN = Pattern.compile("^event:(.+)$", Pattern.MULTILINE); + /** + * Stateful UTF-8 decoder that splits a stream of byte-buffer chunks into complete + * lines. Handles multi-byte characters split across chunk boundaries, and terminates + * a line on {@code "\r\n"}, {@code "\r"} or {@code "\n"} alike, as the SSE wire + * format does. Bytes that do not decode are replaced rather than reported, so a peer + * sending one does not cost the stream. + */ + static final class Utf8LineDecoder { /** - * The sink for emitting parsed response events. + * Undecodable input costs one replacement character rather than the stream: a + * decoder left on the default {@link CodingErrorAction#REPORT} fails the whole + * response over a single byte a peer mangled, and takes with it the lines already + * decoded from the same chunk, because {@link #decode(List)} throws instead of + * returning them. A body cut short mid-character is enough to hit it. This + * matches {@link java.net.http.HttpResponse.BodySubscribers#fromLineSubscriber}, + * the path this decoder replaces, which configured the same two actions. */ - private final FluxSink sink; + private final CharsetDecoder decoder = StandardCharsets.UTF_8.newDecoder() + .onMalformedInput(CodingErrorAction.REPLACE) + .onUnmappableCharacter(CodingErrorAction.REPLACE); - /** - * StringBuilder for accumulating multi-line event data. - */ - private final StringBuilder eventBuilder; + private final CharBuffer charBuffer = CharBuffer.allocate(4096); - /** - * Current event's ID, if specified. - */ - private final AtomicReference currentEventId; + private final StringBuilder leftover = new StringBuilder(); /** - * Current event's type, if specified. + * 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 final AtomicReference currentEventType; + private int scannedForLineTerminator = 0; /** - * The response information from the HTTP response. Send with each event to - * provide context. + * Whether the line just emitted was terminated by a CR, so that a LF opening what + * follows completes that terminator instead of ending a line of its own. A CR is + * emitted on as soon as it arrives, before it is known whether a LF follows it, + * and the two may be split across chunks. */ - private ResponseInfo responseInfo; - - /** - * 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; - - /** - * 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; - } - - @Override - protected void hookOnSubscribe(Subscription subscription) { - - sink.onRequest(n -> { - subscription.request(n); - }); - - // Register disposal callback to cancel subscription when Flux is disposed - sink.onDispose(() -> { - subscription.cancel(); - }); - } - - @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()); - - this.sink.next(new SseResponseEvent(responseInfo, sseEvent)); - this.eventBuilder.setLength(0); + private boolean crTerminatedPreviousLine = false; + + // 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); + + 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); + // Unreachable while the decoder replaces undecodable input, but kept + // so that an error result cannot spin this loop: it is neither an + // underflow nor an overflow. + 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()); } - 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()); + if (result.isUnderflow()) { + if (input.hasRemaining()) { + pendingBytes = ByteBuffer.allocate(input.remaining()); + pendingBytes.put(input).flip(); + } + break; } - upstream().request(1); } - 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); - } - 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)); + } + 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(); + } + catch (CharacterCodingException e) { + throw new RuntimeException(e); } } - } - @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)); + 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(); + leftover.setLength(0); + this.scannedForLineTerminator = 0; + lines.add(last); } - this.sink.complete(); + this.crTerminatedPreviousLine = false; + return lines; } - @Override - protected void hookOnError(Throwable throwable) { - this.sink.error(throwable); + private void drainCharBuffer() { + charBuffer.flip(); + leftover.append(charBuffer); + charBuffer.clear(); } - } - - static class AggregateSubscriber extends BaseSubscriber { - - /** - * The sink for emitting parsed response events. - */ - private final FluxSink sink; - - /** - * StringBuilder for accumulating multi-line event data. - */ - private final StringBuilder eventBuilder; - - /** - * The response information from the HTTP response. Send with each event to - * provide context. - */ - private ResponseInfo responseInfo; - - volatile boolean hasRequestedDemand = false; - - /** - * 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 void extractCompletedLines(List out) { + while (true) { + if (this.crTerminatedPreviousLine) { + if (leftover.length() == 0) { + // The LF, if there is one, is in a chunk that has not arrived. + return; + } + if (leftover.charAt(0) == '\n') { + leftover.delete(0, 1); + } + this.crTerminatedPreviousLine = false; + } + int terminatorIdx = indexOfLineTerminator(this.scannedForLineTerminator); + if (terminatorIdx == -1) { + this.scannedForLineTerminator = leftover.length(); + return; + } + out.add(leftover.substring(0, terminatorIdx)); + this.crTerminatedPreviousLine = leftover.charAt(terminatorIdx) == '\r'; + leftover.delete(0, terminatorIdx + 1); + // What is left starts after the terminator, so none of it has been + // searched yet. + this.scannedForLineTerminator = 0; + } + } /** - * 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 + * Index of the first CR or LF in {@link #leftover} at or after {@code from}, or + * {@code -1} when there is none. */ - public AggregateSubscriber(ResponseInfo responseInfo, FluxSink sink, int maxSize) { - this.sink = sink; - this.eventBuilder = new StringBuilder(); - this.responseInfo = responseInfo; - this.maxSize = maxSize; - } - - @Override - protected void hookOnSubscribe(Subscription subscription) { - - sink.onRequest(n -> { - if (!hasRequestedDemand) { - subscription.request(Long.MAX_VALUE); + private int indexOfLineTerminator(int from) { + for (int i = from; i < leftover.length(); i++) { + char c = leftover.charAt(i); + if (c == '\n' || c == '\r') { + return i; } - 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; } - this.eventBuilder.append(line).append("\n"); + return -1; } - @Override - protected void hookOnComplete() { - - if (hasRequestedDemand) { - String data = this.eventBuilder.toString(); - this.sink.next(new AggregateResponseEvent(responseInfo, data)); - } - - this.sink.complete(); - } + } - @Override - protected void hookOnError(Throwable throwable) { - this.sink.error(throwable); - } + /** + * 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, and a blank line dispatches only when a {@code data:} field was seen, + * whether or not it carried a value. + */ + static final class SseEventParser { - } + private static final Logger logger = LoggerFactory.getLogger(SseEventParser.class); - static class BodilessResponseLineSubscriber extends BaseSubscriber { + private final StringBuilder data = new StringBuilder(); /** - * The sink for emitting parsed response events. + * 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 FluxSink sink; + private final int maxSize; - private final ResponseInfo responseInfo; + private String id; - volatile boolean hasRequestedDemand = false; + private String event; - public BodilessResponseLineSubscriber(ResponseInfo responseInfo, FluxSink sink) { - this.sink = sink; - 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 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)); + SseEvent result = new SseEvent(id, event, data.toString().trim()); + data.setLength(0); + return Optional.of(result); } - this.sink.complete(); + if (line.startsWith("data:")) { + // Every data field appends its value followed by a separator, so a + // valueless `data:` line still marks the event as carrying data and gets + // dispatched with empty data. Servers send such an event to prime a + // stream, and dropping it leaves the request it answers hanging. + String value = line.substring(5).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'); + } + else if (line.startsWith("id:")) { + String rest = line.substring(3); + if (!rest.isEmpty()) { + id = rest.trim(); + } + } + 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 +528,25 @@ 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 line terminator) exceeds {@code maxSize} bytes. + * + *

+ * {@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. * *

- * {@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. + * CR and LF both reset the count, matching the terminators {@link Utf8LineDecoder} + * flushes a line on: whatever empties the decoder's buffer has to refill this budget, + * or a peer framing short lines with CR alone would be aborted for exceeding a bound + * its lines never reach. A CRLF resets twice, which is harmless. */ - 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,9 +559,9 @@ 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; } @@ -581,7 +580,7 @@ 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. + * the whole span added to the running count when the buffer holds none. */ private long lengthOfTrailingRun(ByteBuffer buffer, int position, int limit) { for (int i = limit - 1; i >= position; 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..e163f903d 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,30 @@ 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 lineSubscriberIsResetByLoneCarriageReturns() { + BoundedLineBodySubscriber subscriber = lineSubscriber(); + + // A lone CR terminates a line, so it flushes Utf8LineDecoder's buffer and has to + // refill the budget here too. Otherwise a peer framing short lines with CR alone + // has its response aborted for exceeding a bound its lines never reach. + boolean accepted = true; + for (int i = 0; i < 10 && accepted; i++) { + accepted = subscriber.checkSize(buffer("aa\r")); + } + + assertThat(accepted).isTrue(); + } + @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 +132,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 +144,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 +153,7 @@ void lineSubscriberOnlyCountsFromTheBufferPosition() { @Test void lineSubscriberDoesNotConsumeTheBuffer() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); ByteBuffer buffer = buffer("aaaa\nbbbb"); buffer.position(2); @@ -199,7 +214,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 +226,7 @@ void forwardsBuffersWhileWithinBounds() { @Test void abortsTheResponseWhenTheLineBoundIsExceeded() { - BoundedLineBodySubscriber subscriber = lineSubscriber(); + BoundedLineBodySubscriber subscriber = lineSubscriber(); subscriber.onNext(List.of(buffer("a".repeat(MAX_SIZE + 1)))); @@ -235,7 +250,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 +260,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 +277,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/HttpClientHttpTransportLeakTests.java b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientHttpTransportLeakTests.java new file mode 100644 index 000000000..28e2fb46a --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientHttpTransportLeakTests.java @@ -0,0 +1,107 @@ +/* + * Copyright 2026-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.util.List; +import java.util.Map; +import java.util.function.Function; +import java.util.stream.Stream; + +import io.modelcontextprotocol.spec.McpClientTransport; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.json.gson.GsonMcpJsonMapper; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; +import reactor.test.StepVerifier; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.jupiter.api.Named.named; +import static org.junit.jupiter.params.provider.Arguments.arguments; + +class HttpClientHttpTransportLeakTests { + + 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); + } + + @ParameterizedTest + @MethodSource("httpTransports") + void closeDoesNotRetainOwnedHttpClient(Function httpTransportBuilder) throws Exception { + try (LoopbackMcpHttpServer server = LoopbackMcpHttpServer.start()) { + int selectorThreadsBefore = selectorManagerThreadCount(); + Function, reactor.core.publisher.Mono> handler = Function + .identity(); + + for (int i = 0; i < 12; i++) { + McpClientTransport transport = httpTransportBuilder.apply(server.baseUri().toString()); + + StepVerifier.create(transport.connect(handler)).verifyComplete(); + StepVerifier.create(transport.sendMessage( + new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, "ping", Map.of("iteration", i)))) + .verifyComplete(); + pauseForSelectorStartup(); + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + } + + int selectorThreadsAfter = forceGcUntilStable(); + + assertThat(selectorThreadsAfter) + .describedAs( + "closed transports should not keep owned HttpClient instances alive, remaining threads: %s", + selectorManagerThreadNames()) + .isLessThanOrEqualTo(selectorThreadsBefore + 1); + } + } + + static Stream httpTransports() { + Function streamableHttp = ( + uri) -> HttpClientStreamableHttpTransport.builder(uri).jsonMapper(new GsonMcpJsonMapper()).build(); + Function sse = ( + uri) -> HttpClientSseClientTransport.builder(uri).jsonMapper(new GsonMcpJsonMapper()).build(); + return Stream.of(arguments(named("Streamable HTTP", streamableHttp)), arguments(named("SSE", sse))); + } + +} 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/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..07ed8f1e6 --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/SseEventParserTests.java @@ -0,0 +1,149 @@ +/* + * 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 dataFieldWithEmptyValueStillDispatchesAnEvent() { + // Per the SSE spec a data field appends its value plus a separator, so a lone + // `data:` line leaves the buffer non-empty and the event is dispatched carrying + // empty data. Servers send exactly this to prime a stream, and dropping it leaves + // the request the stream answers hanging. + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed("data:")).isEmpty(); + Optional event = p.feed(""); + assertThat(event).isPresent(); + assertThat(event.get().data()).isEmpty(); + } + + @Test + void dataFieldWithOnlyASpaceIsEquivalentToNoValue() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed("data: ")).isEmpty(); + Optional event = p.feed(""); + assertThat(event).isPresent(); + assertThat(event.get().data()).isEmpty(); + } + + @Test + void blankLineWithNoDataFieldDispatchesNothing() { + // `event:` alone leaves the data buffer empty, which per the spec is not an event + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed("event: message")).isEmpty(); + assertThat(p.feed("")).isEmpty(); + } + + @Test + void valuelessDataFieldIsDispatchedOnFlush() { + SseEventParser p = new SseEventParser(Integer.MAX_VALUE); + assertThat(p.feed("data:")).isEmpty(); + Optional flushed = p.flush(); + assertThat(flushed).isPresent(); + assertThat(flushed.get().data()).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..ce6b3abbe --- /dev/null +++ b/mcp-core/src/test/java/io/modelcontextprotocol/client/transport/Utf8LineDecoderTests.java @@ -0,0 +1,249 @@ +/* + * Copyright 2024-2026 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import java.io.ByteArrayOutputStream; +import java.nio.ByteBuffer; +import java.nio.charset.StandardCharsets; +import java.util.List; + +import io.modelcontextprotocol.client.transport.ResponseSubscribers.Utf8LineDecoder; +import org.junit.jupiter.api.Test; + +import static org.assertj.core.api.Assertions.assertThat; + +class Utf8LineDecoderTests { + + /** + * 0xFF cannot appear anywhere in well-formed UTF-8. One of these is what a peer + * mixing encodings, or a proxy corrupting a byte, puts on the wire. + */ + private static final byte[] INVALID_BYTE = { (byte) 0xFF }; + + /** + * The lead byte of the two-byte sequence for {@code 'é'} (U+00E9, 0xC3 0xA9). + */ + private static final byte[] TRUNCATED_LEAD_BYTE = { (byte) 0xC3 }; + + private static List chunk(String... parts) { + return List.of(toByteBuffers(parts)); + } + + /** + * A chunk whose bytes are passed through verbatim, so that bytes no encoder would + * produce reach the decoder as-is. + */ + private static List rawChunk(byte[]... parts) { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + for (byte[] part : parts) { + out.writeBytes(part); + } + return List.of(ByteBuffer.wrap(out.toByteArray())); + } + + private static byte[] utf8(String text) { + return text.getBytes(StandardCharsets.UTF_8); + } + + 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 trailingCrTerminatesTheLine() { + // A body whose last byte is a CR ends on a terminator, not part-way through a + // line, so there is nothing left to flush. + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("complete\r"))).containsExactly("complete"); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void emptyInput() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(List.of())).isEmpty(); + assertThat(dec.flush()).isEmpty(); + } + + @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 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); + } + + @Test + void loneCrTerminatesLine() { + // SSE takes its line endings from HTML, which terminates on CRLF, CR and LF + // alike, and HttpResponse.BodySubscribers#fromLineSubscriber -- the path this + // decoder replaces -- splits on all three. Splitting on LF alone leaves a + // CR-framed stream as one unterminated run: downstream an "Invalid SSE response + // line", or past BoundedLineBodySubscriber's bound an aborted response. + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("one\rtwo\rthree\r"))).containsExactly("one", "two", "three"); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void blankLinesFramedWithCr() { + // A CR ending a chunk terminates its line, so the CR opening the next one ends an + // empty line rather than completing a CRLF. Values match what + // HttpResponse.BodySubscribers#fromLineSubscriber produces for the same bytes. + assertThat(new Utf8LineDecoder().decode(chunk("\r\r"))).containsExactly("", ""); + + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("one\r"))).containsExactly("one"); + assertThat(dec.decode(chunk("\r"))).containsExactly(""); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void sseFramedWithCrOnlyIsSplitIntoFieldLines() { + // The same stream as the SSE parser downstream has to receive it: one line per + // field, and the empty line that ends the event. + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("event: message\rdata: {\"a\":1}\r\r"))).containsExactly("event: message", + "data: {\"a\":1}", ""); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void crLfSplitAcrossChunks() { + // Splitting on a lone CR means emitting the line as soon as the CR arrives, so a + // LF opening the next chunk is the tail of a CRLF rather than an empty line of + // its own. The terminator also sits exactly at the point the previous search for + // one stopped. + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(chunk("hello\r"))).containsExactly("hello"); + assertThat(dec.decode(chunk("\nworld\r\n"))).containsExactly("world"); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void malformedByteIsReplacedAndOtherLinesArePreserved() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(rawChunk(utf8("one\ncaf"), INVALID_BYTE, utf8("e\nthree\n")))).containsExactly("one", + "caf\uFFFDe", "three"); + assertThat(dec.flush()).isEmpty(); + } + + @Test + void trailingTruncatedCharacterIsReplacedOnFlush() { + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(rawChunk(utf8("one\ncaf"), TRUNCATED_LEAD_BYTE))).containsExactly("one"); + assertThat(dec.flush()).containsExactly("caf\uFFFD"); + } + + @Test + void incompleteMultiByteSequenceFollowedByValidDataIsReplaced() { + // character is cut short by a chunk boundary + // "€" is U+20AC → 0xE2 0x82 0xAC in UTF-8; only the first two bytes arrive. + byte[] euro = "€".getBytes(StandardCharsets.UTF_8); + + Utf8LineDecoder dec = new Utf8LineDecoder(); + assertThat(dec.decode(rawChunk(new byte[] { euro[0], euro[1] }))).isEmpty(); + assertThat(dec.decode(chunk("x\n"))).containsExactly("\uFFFDx"); + } + +} 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/HttpClientStreamableHttpTransportEmptyJsonResponseTest.java b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportEmptyResponseTests.java similarity index 54% rename from mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportEmptyJsonResponseTest.java rename to mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportEmptyResponseTests.java index c2d19ef67..6668941b6 100644 --- a/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportEmptyJsonResponseTest.java +++ b/mcp-test/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportEmptyResponseTests.java @@ -11,9 +11,12 @@ import static org.mockito.Mockito.verify; import java.io.IOException; +import java.io.OutputStream; import java.net.InetSocketAddress; import java.net.URI; import java.net.URISyntaxException; +import java.nio.charset.StandardCharsets; +import java.util.Map; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; @@ -29,11 +32,12 @@ import reactor.test.StepVerifier; /** - * Handles emplty application/json response with 200 OK status code. + * Handles 200 OK responses that carry no usable body, either as an empty application/json + * document or as a text/event-stream containing nothing but a stream primer. * * @author codezkk */ -public class HttpClientStreamableHttpTransportEmptyJsonResponseTest { +public class HttpClientStreamableHttpTransportEmptyResponseTests { static int PORT = TomcatTestUtil.findAvailablePort(); @@ -41,6 +45,20 @@ public class HttpClientStreamableHttpTransportEmptyJsonResponseTest { static HttpServer server; + /** + * An SSE event with an {@code event:} field but no data, which some servers send to + * open the response stream before any JSON-RPC payload is available. Note the + * valueless {@code data:} field: per the SSE spec this is identical to {@code data: } + * with a trailing space. + * @see SEP-1699 + */ + private static final byte[] SSE_PRIMER = """ + event: message + data: + + """.getBytes(StandardCharsets.UTF_8); + @BeforeAll static void startContainer() throws IOException { @@ -53,6 +71,24 @@ static void startContainer() throws IOException { exchange.close(); }); + // 200 OK text/event-stream carrying only a primer, for POSTs. The + // server-initiated GET stream is refused so that the transport falls back to + // request-response mode and the POST is the only thing under test. + server.createContext("/mcp-sse-primer", exchange -> { + try (exchange) { + if (!"POST".equals(exchange.getRequestMethod())) { + exchange.sendResponseHeaders(405, -1); + return; + } + exchange.getRequestBody().readAllBytes(); + exchange.getResponseHeaders().set("Content-Type", "text/event-stream"); + exchange.sendResponseHeaders(200, SSE_PRIMER.length); + try (OutputStream out = exchange.getResponseBody()) { + out.write(SSE_PRIMER); + } + } + }); + server.setExecutor(null); server.start(); } @@ -91,4 +127,25 @@ void testNotificationInitialized() throws URISyntaxException { } + /** + * A POST answered with {@code 200 text/event-stream} whose body holds only a stream + * primer must still complete, because the primer tells the client the stream is live + * and the message has been accepted. The primer's {@code data:} field carries no + * value, so this only holds as long as such a field still produces an event: a parser + * that drops it leaves no event to fire the transport's first-message callback, and + * {@code sendMessage} then never completes at all. + */ + @Test + @Timeout(5) + void testNotificationAnsweredWithSsePrimerOnly() { + + var transport = HttpClientStreamableHttpTransport.builder(host).endpoint("/mcp-sse-primer").build(); + + var testMessage = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION, "notifications/initialized", + Map.of()); + + StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete(); + + } + } 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 + "\"}}"; + } + +}