refactor: improve SSE client transpor testing

- Extract eventStream() method
- Enhance connect() method with improved SSE event processing and error handling
- Introduce MyResponseInfo record and enhanced TestHttpClientSseClientTransport
- Add collections to track JSON-RPC requests, notifications, and responses

Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>

Author: tzolov

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

@@ -321,9 +321,89 @@ public HttpClientSseClientTransport build() {
 
 	}
 
+	protected Flux<ResponseSubscribers.SseResponseEvent> eventStream() {
+
+		HttpRequest request = requestBuilder.copy()
+			.uri(Utils.resolveUri(this.baseUri, this.sseEndpoint))
+			.header("Accept", "text/event-stream")
+			.header("Cache-Control", "no-cache")
+			.GET()
+			.build();
+
+		Flux<ResponseSubscribers.SseResponseEvent> bla = Flux.<ResponseEvent>create(sseSink -> this.httpClient
+			.sendAsync(request, responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink))
+			.exceptionallyCompose(e -> {
+				sseSink.error(e);
+				return CompletableFuture.failedFuture(e);
+			})).map(responseEvent -> (ResponseSubscribers.SseResponseEvent) responseEvent);
+		return bla;
+	}
+
 	@Override
 	public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
 
+		return Mono.create(sink -> {
+
+			Flux<ResponseSubscribers.SseResponseEvent> events = eventStream();
+
+			Disposable connection = events.flatMap(responseEvent -> {
+				if (isClosing) {
+					return Mono.empty();
+				}
+
+				int statusCode = responseEvent.responseInfo().statusCode();
+
+				if (statusCode >= 200 && statusCode < 300) {
+					try {
+						if (ENDPOINT_EVENT_TYPE.equals(responseEvent.sseEvent().event())) {
+							String messageEndpointUri = responseEvent.sseEvent().data();
+							if (this.messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) {
+								sink.success();
+								return Flux.empty(); // No further processing needed
+							}
+							else {
+								sink.error(new McpError("Failed to handle SSE endpoint event"));
+							}
+						}
+						else if (MESSAGE_EVENT_TYPE.equals(responseEvent.sseEvent().event())) {
+							JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper,
+									responseEvent.sseEvent().data());
+							sink.success();
+							return Flux.just(message);
+						}
+						else {
+							logger.error("Received unrecognized SSE event type: {}", responseEvent.sseEvent().event());
+							sink.error(new McpError(
+									"Received unrecognized SSE event type: " + responseEvent.sseEvent().event()));
+						}
+					}
+					catch (IOException e) {
+						logger.error("Error processing SSE event", e);
+						sink.error(new McpError("Error processing SSE event"));
+					}
+				}
+				return Flux.<McpSchema.JSONRPCMessage>error(
+						new RuntimeException("Failed to send message: " + responseEvent));
+
+			}).flatMap(jsonRpcMessage -> handler.apply(Mono.just(jsonRpcMessage))).onErrorComplete(t -> {
+				if (!isClosing) {
+					logger.warn("SSE stream observed an error", t);
+					sink.error(t);
+				}
+				return true;
+			}).doFinally(s -> {
+				Disposable ref = this.sseSubscription.getAndSet(null);
+				if (ref != null && !ref.isDisposed()) {
+					ref.dispose();
+				}
+			}).contextWrite(sink.contextView()).subscribe();
+
+			this.sseSubscription.set(connection);
+		});
+	}
+
+	public Mono<Void> connect2(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
+
 		return Mono.create(sink -> {
 
 			HttpRequest request = requestBuilder.copy()
@@ -333,6 +413,13 @@ public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> h
 				.GET()
 				.build();
 
+			Flux<ResponseSubscribers.SseResponseEvent> bla = Flux.<ResponseEvent>create(sseSink -> this.httpClient
+				.sendAsync(request, responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink))
+				.exceptionallyCompose(e -> {
+					sseSink.error(e);
+					return CompletableFuture.failedFuture(e);
+				})).map(responseEvent -> (ResponseSubscribers.SseResponseEvent) responseEvent);
+
 			Disposable connection = Flux.<ResponseEvent>create(sseSink -> this.httpClient
 				.sendAsync(request, responseInfo -> ResponseSubscribers.sseToBodySubscriber(responseInfo, sseSink))
 				.exceptionallyCompose(e -> {

mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransportTests.java

@@ -4,34 +4,43 @@
 
 package io.modelcontextprotocol.client.transport;
 
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+
 import java.net.URI;
 import java.net.http.HttpClient;
+import java.net.http.HttpClient.Version;
+import java.net.http.HttpHeaders;
 import java.net.http.HttpRequest;
+import java.net.http.HttpResponse.ResponseInfo;
 import java.time.Duration;
 import java.util.Map;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
-import com.fasterxml.jackson.databind.ObjectMapper;
-import io.modelcontextprotocol.spec.McpSchema;
-import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;
 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.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.wait.strategy.Wait;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+import io.modelcontextprotocol.client.transport.ResponseSubscribers.SseResponseEvent;
+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 reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 import reactor.core.publisher.Sinks;
 import reactor.test.StepVerifier;
 
-import org.springframework.http.codec.ServerSentEvent;
-
-import static org.assertj.core.api.Assertions.assertThat;
-import static org.assertj.core.api.Assertions.assertThatCode;
-
 /**
  * Tests for the {@link HttpClientSseClientTransport} class.
  *
@@ -51,28 +60,70 @@ class HttpClientSseClientTransportTests {
 
 	private TestHttpClientSseClientTransport transport;
 
+	public record MyResponseInfo(int statusCode, HttpHeaders headers, Version version) implements ResponseInfo {
+		MyResponseInfo(int statusCode, HttpHeaders headers) {
+			this(statusCode, headers, Version.HTTP_1_1);
+		}
+
+		MyResponseInfo(int statusCode) {
+			this(statusCode, HttpHeaders.of(Map.of(), (k, v) -> true), Version.HTTP_1_1);
+		}
+	}
+
 	// Test class to access protected methods
 	static class TestHttpClientSseClientTransport extends HttpClientSseClientTransport {
 
 		private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
 
-		private Sinks.Many<ServerSentEvent<String>> events = Sinks.many().unicast().onBackpressureBuffer();
+		private Sinks.Many<SseResponseEvent> events = Sinks.many().unicast().onBackpressureBuffer();
 
 		public TestHttpClientSseClientTransport(final String baseUri) {
-			super(HttpClient.newHttpClient(), HttpRequest.newBuilder(), baseUri, "/sse", new ObjectMapper());
+			super(HttpClient.newBuilder().version(HttpClient.Version.HTTP_1_1).build(),
+					HttpRequest.newBuilder().header("Content-Type", "application/json"), baseUri, "/sse",
+					new ObjectMapper());
+		}
+
+		CopyOnWriteArrayList<JSONRPCRequest> requestMessages = new CopyOnWriteArrayList<>();
+
+		CopyOnWriteArrayList<JSONRPCNotification> notificationMessages = new CopyOnWriteArrayList<>();
+
+		CopyOnWriteArrayList<JSONRPCResponse> responseMessages = new CopyOnWriteArrayList<>();
+
+		Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler = (messageMono) -> messageMono
+			.doOnNext(message -> {
+				// System.out.println("Received message $$$$$$$$$$$$$$: " + message);
+				if (message instanceof JSONRPCRequest request) {
+					requestMessages.add(request);
+				}
+				else if (message instanceof JSONRPCNotification notificaiton) {
+					notificationMessages.add(notificaiton);
+				}
+				else if (message instanceof JSONRPCResponse response) {
+					responseMessages.add(response);
+				}
+				else {
+					throw new IllegalArgumentException("Unsupported message type: " + message.getClass());
+				}
+			});
+
+		@Override
+		protected Flux<SseResponseEvent> eventStream() {
+			return super.eventStream().mergeWith(events.asFlux());
 		}
 
 		public int getInboundMessageCount() {
 			return inboundMessageCount.get();
 		}
 
 		public void simulateEndpointEvent(String jsonMessage) {
-			events.tryEmitNext(ServerSentEvent.<String>builder().event("endpoint").data(jsonMessage).build());
+			events.tryEmitNext(new SseResponseEvent(new MyResponseInfo(200),
+					new ResponseSubscribers.SseEvent(null, "endpoint", jsonMessage)));
 			inboundMessageCount.incrementAndGet();
 		}
 
 		public void simulateMessageEvent(String jsonMessage) {
-			events.tryEmitNext(ServerSentEvent.<String>builder().event("message").data(jsonMessage).build());
+			events.tryEmitNext(new SseResponseEvent(new MyResponseInfo(200),
+					new ResponseSubscribers.SseEvent(null, "message", jsonMessage)));
 			inboundMessageCount.incrementAndGet();
 		}
 
@@ -88,7 +139,7 @@ void startContainer() {
 	void setUp() {
 		startContainer();
 		transport = new TestHttpClientSseClientTransport(host);
-		transport.connect(Function.identity()).block();
+		transport.connect(transport.handler).block();
 	}
 
 	@AfterEach
@@ -123,6 +174,7 @@ void testMessageProcessing() {
 		StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
 
 		assertThat(transport.getInboundMessageCount()).isEqualTo(1);
+		assertThat(transport.requestMessages).hasSize(1);
 	}
 
 	@Test