Address connect hangs in legacy SSE connect and streamline code paths

Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>

Author: chemicL

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

@@ -11,7 +11,7 @@
 import java.net.http.HttpResponse;
 import java.time.Duration;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -389,14 +389,6 @@ public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> h
 			var transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY);
 			return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, null, transportContext));
 		}).flatMap(requestBuilder -> Mono.create(sink -> {
-			// Once connect() has completed, a later failure can no longer be reported
-			// through its sink: signalling it there would only have Reactor drop it.
-			AtomicBoolean connected = new AtomicBoolean();
-			Runnable markConnected = () -> {
-				if (connected.compareAndSet(false, true)) {
-					sink.success();
-				}
-			};
 			Disposable connection = ResponseBodyHandlers.sendAsync(this.httpClient, requestBuilder.build())
 				.flatMapMany(response -> {
 					if (isClosing) {
@@ -417,7 +409,10 @@ public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> h
 								"Failed to connect to SSE stream: " + statusCode);
 					}
 				})
-				.flatMap(sseEvent -> {
+				// Every successfully processed event yields exactly one element, empty
+				// when it carries no message, so that the first one can mark the
+				// connection as established.
+				.<Optional<JSONRPCMessage>>handle((sseEvent, events) -> {
 					try {
 						if (ENDPOINT_EVENT_TYPE.equals(sseEvent.event())) {
 							String messageEndpointUri = sseEvent.data();
@@ -426,46 +421,61 @@ public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> h
 							}
 							catch (InvalidSseMessageEndpointException e) {
 								this.messageEndpointSink.tryEmitError(e);
-								return Flux.error(e);
+								events.error(e);
+								return;
 							}
 							if (this.messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) {
-								markConnected.run();
-								return Flux.empty(); // No further processing needed
+								events.next(Optional.empty());
+							}
+							else {
+								events.error(new McpTransportException("Failed to handle SSE endpoint event"));
 							}
-							return Flux.error(new McpTransportException("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)");
-								markConnected.run();
-								return Flux.<McpSchema.JSONRPCMessage>empty();
+								events.next(Optional.empty());
+							}
+							else {
+								events.next(Optional.of(McpSchema.deserializeJsonRpcMessage(jsonMapper, data)));
 							}
-							JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(jsonMapper, data);
-							markConnected.run();
-							return Flux.just(message);
 						}
 						else {
 							logger.debug("Received unrecognized SSE event type: {}", sseEvent);
-							markConnected.run();
-							return Flux.<McpSchema.JSONRPCMessage>empty();
+							events.next(Optional.empty());
 						}
 					}
 					catch (IOException e) {
-						return Flux.<McpSchema.JSONRPCMessage>error(
-								new McpTransportException("Error processing SSE event", e));
+						events.error(new McpTransportException("Error processing SSE event", e));
+					}
+				})
+				// connect() is resolved by the first signal only: any later failure is
+				// merely logged below, as connect() has already completed by then.
+				.switchOnFirst((first, events) -> {
+					if (first.hasValue()) {
+						sink.success();
+					}
+					else if (first.isOnError()) {
+						sink.error(first.getThrowable());
 					}
+					else if (first.isOnComplete()) {
+						sink.error(new McpTransportException("SSE stream closed before any event was received"));
+					}
+					return events;
 				})
-				.flatMap(jsonRpcMessage -> handler.apply(Mono.just(jsonRpcMessage)))
+				.<JSONRPCMessage>handle((message, messages) -> message.ifPresent(messages::next))
+				.flatMap(message -> handler.apply(Mono.just(message)))
 				.onErrorComplete(t -> {
 					if (!isClosing) {
 						logger.warn("SSE stream observed an error", t);
-						if (connected.compareAndSet(false, true)) {
-							sink.error(t);
-						}
 					}
 					return true;
 				})
+				// A closeGracefully() before the first signal cancels the stream:
+				// complete
+				// connect() instead of leaving it pending. A no-op once it has resolved.
+				.doOnCancel(sink::success)
 				.doFinally(s -> {
 					Disposable ref = this.sseSubscription.getAndSet(null);
 					if (ref != null && !ref.isDisposed()) {

mcp-core/src/test/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransportConnectTests.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2026-2026 the original author or authors.
+ */
+
+package io.modelcontextprotocol.client.transport;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.time.Duration;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+
+import com.sun.net.httpserver.HttpHandler;
+import com.sun.net.httpserver.HttpServer;
+import io.modelcontextprotocol.spec.McpTransportException;
+import io.modelcontextprotocol.spec.json.gson.GsonMcpJsonMapper;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import reactor.test.StepVerifier;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Verifies that {@link HttpClientSseClientTransport#connect} always resolves, even when
+ * the SSE stream ends, fails or is closed before its first event.
+ */
+class HttpClientSseClientTransportConnectTests {
+
+	// Only bounds a regression: every test resolves without waiting on it.
+	private static final Duration TIMEOUT = Duration.ofSeconds(5);
+
+	private final ExecutorService executor = Executors.newCachedThreadPool();
+
+	private final CountDownLatch releaseResponse = new CountDownLatch(1);
+
+	private HttpServer server;
+
+	@AfterEach
+	void tearDown() {
+		this.releaseResponse.countDown();
+		if (this.server != null) {
+			this.server.stop(0);
+		}
+		this.executor.shutdownNow();
+	}
+
+	@Test
+	void connectFailsWhenStreamEndsBeforeAnyEvent() throws IOException {
+		HttpClientSseClientTransport transport = transport(exchange -> {
+			exchange.getResponseHeaders().add("Content-Type", "text/event-stream");
+			exchange.sendResponseHeaders(200, -1);
+			exchange.close();
+		});
+
+		StepVerifier.create(transport.connect(Function.identity()))
+			.expectErrorSatisfies(e -> assertThat(e).isInstanceOf(McpTransportException.class)
+				.hasMessageContaining("before any event"))
+			.verify(TIMEOUT);
+	}
+
+	@Test
+	void connectFailsWhenStreamErrorsBeforeAnyEventWhileClosing() throws IOException {
+		HttpClientSseClientTransport transport = transport(exchange -> {
+			// The server drops the connection without responding.
+			throw new IOException("dropped");
+		});
+		transport.closeGracefully().block(TIMEOUT);
+
+		StepVerifier.create(transport.connect(Function.identity())).expectError().verify(TIMEOUT);
+	}
+
+	@Test
+	void connectCompletesWhenClosedBeforeAnyEvent() throws IOException {
+		CountDownLatch streamOpened = new CountDownLatch(1);
+		HttpClientSseClientTransport transport = transport(exchange -> {
+			exchange.getResponseHeaders().add("Content-Type", "text/event-stream");
+			exchange.sendResponseHeaders(200, 0);
+			exchange.getResponseBody().flush();
+			streamOpened.countDown();
+			try {
+				this.releaseResponse.await();
+			}
+			catch (InterruptedException e) {
+				Thread.currentThread().interrupt();
+			}
+			exchange.close();
+		});
+
+		StepVerifier.create(transport.connect(Function.identity())).then(() -> {
+			try {
+				assertThat(streamOpened.await(TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)).isTrue();
+			}
+			catch (InterruptedException e) {
+				throw new IllegalStateException(e);
+			}
+			transport.closeGracefully().block(TIMEOUT);
+		}).expectComplete().verify(TIMEOUT);
+	}
+
+	private HttpClientSseClientTransport transport(HttpHandler sseHandler) throws IOException {
+		this.server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+		this.server.setExecutor(this.executor);
+		this.server.createContext("/sse", sseHandler);
+		this.server.start();
+		return HttpClientSseClientTransport.builder("http://127.0.0.1:" + this.server.getAddress().getPort())
+			.jsonMapper(new GsonMcpJsonMapper())
+			.build();
+	}
+
+}