fix: improve streamable HTTP session reinitialization (#459)
Implements the MCP spec guidelines for streamable HTTP (re)initialization: - Server MAY terminate session and MUST respond with HTTP 404 for terminated session IDs - Client MUST start new session when receiving HTTP 404 for requests with session ID Changes: - Replace generic McpError with McpTransportException for transport-layer errors - Only throw McpTransportSessionNotFoundException when session ID is present in request (per spec: 404 with session ID means session terminated, without means general error) - Enhance error messages with more context (status codes, response events) - Use RuntimeException for non-transport specific SSE endpoint failures - Ensure consistent error handling across HTTP client transports - Improve error handling with standard Java exceptions. Replace generic McpError with appropriate standard exceptions: - Use IllegalArgumentException for invalid input parameters - Use IllegalStateException for state-related issues - Use RuntimeException wrapper for initialization failures - Use McpError.builder() with proper error codes for protocol errors Fixes #459 Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
Christian Tzolov
parent
ed2c35ea02
commit
ea50838d06
9 changed files with 955 additions and 90 deletions
|
|
@ -31,6 +31,7 @@ import io.modelcontextprotocol.spec.HttpHeaders;
|
|||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportException;
|
||||
import io.modelcontextprotocol.spec.McpTransportSession;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import io.modelcontextprotocol.spec.McpTransportStream;
|
||||
|
|
@ -70,6 +71,8 @@ import reactor.util.function.Tuples;
|
|||
*/
|
||||
public class WebClientStreamableHttpTransport implements McpClientTransport {
|
||||
|
||||
private static final String MISSING_SESSION_ID = "[missing_session_id]";
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(WebClientStreamableHttpTransport.class);
|
||||
|
||||
private static final String MCP_PROTOCOL_VERSION = ProtocolVersions.MCP_2025_03_26;
|
||||
|
|
@ -221,8 +224,13 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
|
|||
return Flux.empty();
|
||||
}
|
||||
else if (isNotFound(response)) {
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
return mcpSessionNotFoundError(sessionIdRepresentation);
|
||||
if (transportSession.sessionId().isPresent()) {
|
||||
String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession);
|
||||
return mcpSessionNotFoundError(sessionIdRepresentation);
|
||||
}
|
||||
else {
|
||||
return this.extractError(response, MISSING_SESSION_ID);
|
||||
}
|
||||
}
|
||||
else {
|
||||
return response.<McpSchema.JSONRPCMessage>createError().doOnError(e -> {
|
||||
|
|
@ -318,10 +326,10 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
|
|||
}
|
||||
}
|
||||
else {
|
||||
if (isNotFound(response)) {
|
||||
if (isNotFound(response) && !sessionRepresentation.equals(MISSING_SESSION_ID)) {
|
||||
return mcpSessionNotFoundError(sessionRepresentation);
|
||||
}
|
||||
return extractError(response, sessionRepresentation);
|
||||
return this.extractError(response, sessionRepresentation);
|
||||
}
|
||||
})
|
||||
.flatMap(jsonRpcMessage -> this.handler.get().apply(Mono.just(jsonRpcMessage)))
|
||||
|
|
@ -362,10 +370,10 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
|
|||
McpSchema.JSONRPCResponse.class);
|
||||
jsonRpcError = jsonRpcResponse.error();
|
||||
toPropagate = jsonRpcError != null ? new McpError(jsonRpcError)
|
||||
: new McpError("Can't parse the jsonResponse " + jsonRpcResponse);
|
||||
: new McpTransportException("Can't parse the jsonResponse " + jsonRpcResponse);
|
||||
}
|
||||
catch (IOException ex) {
|
||||
toPropagate = new RuntimeException("Sending request failed", e);
|
||||
toPropagate = new McpTransportException("Sending request failed, " + e.getMessage(), e);
|
||||
logger.debug("Received content together with {} HTTP code response: {}", response.statusCode(), body);
|
||||
}
|
||||
|
||||
|
|
@ -374,7 +382,11 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
|
|||
// invalidate the session
|
||||
// https://github.com/modelcontextprotocol/typescript-sdk/issues/389
|
||||
if (responseException.getStatusCode().isSameCodeAs(HttpStatus.BAD_REQUEST)) {
|
||||
return Mono.error(new McpTransportSessionNotFoundException(sessionRepresentation, toPropagate));
|
||||
if (!sessionRepresentation.equals(MISSING_SESSION_ID)) {
|
||||
return Mono.error(new McpTransportSessionNotFoundException(sessionRepresentation, toPropagate));
|
||||
}
|
||||
return Mono.error(new McpTransportException("Received 400 BAD REQUEST for session "
|
||||
+ sessionRepresentation + ". " + toPropagate.getMessage(), toPropagate));
|
||||
}
|
||||
return Mono.error(toPropagate);
|
||||
}).flux();
|
||||
|
|
@ -403,7 +415,7 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
|
|||
}
|
||||
|
||||
private static String sessionIdOrPlaceholder(McpTransportSession<?> transportSession) {
|
||||
return transportSession.sessionId().orElse("[missing_session_id]");
|
||||
return transportSession.sessionId().orElse(MISSING_SESSION_ID);
|
||||
}
|
||||
|
||||
private Flux<McpSchema.JSONRPCMessage> directResponseFlux(McpSchema.JSONRPCMessage sentMessage,
|
||||
|
|
@ -421,8 +433,7 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
|
|||
}
|
||||
}
|
||||
catch (IOException e) {
|
||||
// TODO: this should be a McpTransportError
|
||||
s.error(e);
|
||||
s.error(new McpTransportException(e));
|
||||
}
|
||||
}).flatMapIterable(Function.identity());
|
||||
}
|
||||
|
|
@ -449,7 +460,7 @@ public class WebClientStreamableHttpTransport implements McpClientTransport {
|
|||
return Tuples.of(Optional.ofNullable(event.id()), List.of(message));
|
||||
}
|
||||
catch (IOException ioException) {
|
||||
throw new McpError("Error parsing JSON-RPC message: " + event.data());
|
||||
throw new McpTransportException("Error parsing JSON-RPC message: " + event.data(), ioException);
|
||||
}
|
||||
}
|
||||
else {
|
||||
|
|
|
|||
|
|
@ -14,7 +14,6 @@ import com.fasterxml.jackson.databind.ObjectMapper;
|
|||
|
||||
import io.modelcontextprotocol.spec.HttpHeaders;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage;
|
||||
import io.modelcontextprotocol.spec.ProtocolVersions;
|
||||
|
|
@ -197,8 +196,6 @@ public class WebFluxSseClientTransport implements McpClientTransport {
|
|||
* @param handler a function that processes incoming JSON-RPC messages and returns
|
||||
* responses
|
||||
* @return a Mono that completes when the connection is fully established
|
||||
* @throws McpError if there's an error processing SSE events or if an unrecognized
|
||||
* event type is received
|
||||
*/
|
||||
@Override
|
||||
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
|
||||
|
|
@ -215,7 +212,7 @@ public class WebFluxSseClientTransport implements McpClientTransport {
|
|||
else {
|
||||
// TODO: clarify with the spec if multiple events can be
|
||||
// received
|
||||
s.error(new McpError("Failed to handle SSE endpoint event"));
|
||||
s.error(new RuntimeException("Failed to handle SSE endpoint event"));
|
||||
}
|
||||
}
|
||||
else if (MESSAGE_EVENT_TYPE.equals(event.event())) {
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import java.time.Duration;
|
|||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
|
@ -26,6 +27,7 @@ import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvide
|
|||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
@Timeout(15)
|
||||
class WebFluxSseIntegrationTests extends AbstractMcpClientServerIntegrationTests {
|
||||
|
||||
private static final int PORT = TestUtil.findAvailablePort();
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import java.time.Duration;
|
|||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
|
@ -26,6 +27,7 @@ import io.modelcontextprotocol.server.transport.WebFluxStatelessServerTransport;
|
|||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
@Timeout(15)
|
||||
class WebFluxStatelessIntegrationTests extends AbstractStatelessIntegrationTests {
|
||||
|
||||
private static final int PORT = TestUtil.findAvailablePort();
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import java.time.Duration;
|
|||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
|
@ -26,6 +27,7 @@ import io.modelcontextprotocol.server.transport.WebFluxStreamableServerTransport
|
|||
import reactor.netty.DisposableServer;
|
||||
import reactor.netty.http.server.HttpServer;
|
||||
|
||||
@Timeout(15)
|
||||
class WebFluxStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests {
|
||||
|
||||
private static final int PORT = TestUtil.findAvailablePort();
|
||||
|
|
|
|||
|
|
@ -0,0 +1,404 @@
|
|||
/*
|
||||
* Copyright 2025-2025 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client.transport;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.timeout;
|
||||
import static org.mockito.Mockito.verify;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.net.InetSocketAddress;
|
||||
import java.time.Duration;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Consumer;
|
||||
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Disabled;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
||||
import com.sun.net.httpserver.HttpServer;
|
||||
|
||||
import io.modelcontextprotocol.server.TestUtil;
|
||||
import io.modelcontextprotocol.spec.HttpHeaders;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportException;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import io.modelcontextprotocol.spec.ProtocolVersions;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
/**
|
||||
* Tests for error handling in WebClientStreamableHttpTransport. Addresses concurrency
|
||||
* issues with proper Reactor patterns.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
@Timeout(15)
|
||||
public class WebClientStreamableHttpTransportErrorHandlingTest {
|
||||
|
||||
private static final int PORT = TestUtil.findAvailablePort();
|
||||
|
||||
private static final String HOST = "http://localhost:" + PORT;
|
||||
|
||||
private HttpServer server;
|
||||
|
||||
private AtomicReference<Integer> serverResponseStatus = new AtomicReference<>(200);
|
||||
|
||||
private AtomicReference<String> currentServerSessionId = new AtomicReference<>(null);
|
||||
|
||||
private AtomicReference<String> lastReceivedSessionId = new AtomicReference<>(null);
|
||||
|
||||
private McpClientTransport transport;
|
||||
|
||||
// Initialize latches for proper request synchronization
|
||||
CountDownLatch firstRequestLatch;
|
||||
|
||||
CountDownLatch secondRequestLatch;
|
||||
|
||||
CountDownLatch getRequestLatch;
|
||||
|
||||
@BeforeEach
|
||||
void startServer() throws IOException {
|
||||
|
||||
// Initialize latches for proper synchronization
|
||||
firstRequestLatch = new CountDownLatch(1);
|
||||
secondRequestLatch = new CountDownLatch(1);
|
||||
getRequestLatch = new CountDownLatch(1);
|
||||
|
||||
server = HttpServer.create(new InetSocketAddress(PORT), 0);
|
||||
|
||||
// Configure the /mcp endpoint with dynamic response
|
||||
server.createContext("/mcp", exchange -> {
|
||||
String method = exchange.getRequestMethod();
|
||||
|
||||
if ("GET".equals(method)) {
|
||||
// This is the SSE connection attempt after session establishment
|
||||
getRequestLatch.countDown();
|
||||
// Return 405 Method Not Allowed to indicate SSE not supported
|
||||
exchange.sendResponseHeaders(405, 0);
|
||||
exchange.close();
|
||||
return;
|
||||
}
|
||||
|
||||
String requestSessionId = exchange.getRequestHeaders().getFirst(HttpHeaders.MCP_SESSION_ID);
|
||||
lastReceivedSessionId.set(requestSessionId);
|
||||
|
||||
int status = serverResponseStatus.get();
|
||||
|
||||
// Track which request this is
|
||||
if (firstRequestLatch.getCount() > 0) {
|
||||
// // First request - should have no session ID
|
||||
firstRequestLatch.countDown();
|
||||
}
|
||||
else if (secondRequestLatch.getCount() > 0) {
|
||||
// Second request - should have session ID
|
||||
secondRequestLatch.countDown();
|
||||
}
|
||||
|
||||
exchange.getResponseHeaders().set("Content-Type", "application/json");
|
||||
|
||||
// Don't include session ID in 404 and 400 responses - the implementation
|
||||
// checks if the transport has a session stored locally
|
||||
String responseSessionId = currentServerSessionId.get();
|
||||
if (responseSessionId != null && status == 200) {
|
||||
exchange.getResponseHeaders().set(HttpHeaders.MCP_SESSION_ID, responseSessionId);
|
||||
}
|
||||
if (status == 200) {
|
||||
String response = "{\"jsonrpc\":\"2.0\",\"result\":{},\"id\":\"test-id\"}";
|
||||
exchange.sendResponseHeaders(200, response.length());
|
||||
exchange.getResponseBody().write(response.getBytes());
|
||||
}
|
||||
else {
|
||||
exchange.sendResponseHeaders(status, 0);
|
||||
}
|
||||
exchange.close();
|
||||
});
|
||||
|
||||
server.setExecutor(null);
|
||||
server.start();
|
||||
|
||||
transport = WebClientStreamableHttpTransport.builder(WebClient.builder().baseUrl(HOST)).build();
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
void stopServer() {
|
||||
if (server != null) {
|
||||
server.stop(0);
|
||||
}
|
||||
StepVerifier.create(transport.closeGracefully()).verifyComplete();
|
||||
}
|
||||
|
||||
/**
|
||||
* Test that 404 response WITHOUT session ID throws McpTransportException (not
|
||||
* SessionNotFoundException)
|
||||
*/
|
||||
@Test
|
||||
void test404WithoutSessionId() {
|
||||
serverResponseStatus.set(404);
|
||||
currentServerSessionId.set(null); // No session ID in response
|
||||
|
||||
var testMessage = createTestMessage();
|
||||
|
||||
StepVerifier.create(transport.sendMessage(testMessage))
|
||||
.expectErrorMatches(throwable -> throwable instanceof McpTransportException
|
||||
&& throwable.getMessage().contains("Not Found") && throwable.getMessage().contains("404")
|
||||
&& !(throwable instanceof McpTransportSessionNotFoundException))
|
||||
.verify(Duration.ofSeconds(5));
|
||||
}
|
||||
|
||||
/**
|
||||
* Test that 404 response WITH session ID throws McpTransportSessionNotFoundException
|
||||
* Fixed version using proper async coordination
|
||||
*/
|
||||
@Test
|
||||
void test404WithSessionId() throws InterruptedException {
|
||||
// First establish a session
|
||||
serverResponseStatus.set(200);
|
||||
currentServerSessionId.set("test-session-123");
|
||||
|
||||
// Set up exception handler to verify session invalidation
|
||||
@SuppressWarnings("unchecked")
|
||||
Consumer<Throwable> exceptionHandler = mock(Consumer.class);
|
||||
transport.setExceptionHandler(exceptionHandler);
|
||||
|
||||
// Connect with handler
|
||||
StepVerifier.create(transport.connect(msg -> msg)).verifyComplete();
|
||||
|
||||
// Send initial message to establish session
|
||||
var testMessage = createTestMessage();
|
||||
|
||||
// Send first message to establish session
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Wait for first request to complete
|
||||
assertThat(firstRequestLatch.await(5, TimeUnit.SECONDS)).isTrue();
|
||||
|
||||
// Wait for the GET request (SSE connection attempt) to complete
|
||||
assertThat(getRequestLatch.await(5, TimeUnit.SECONDS)).isTrue();
|
||||
|
||||
// Now return 404 for next request
|
||||
serverResponseStatus.set(404);
|
||||
|
||||
// Use delaySubscription to ensure session is fully processed before next
|
||||
// request
|
||||
StepVerifier.create(Mono.delay(Duration.ofMillis(200)).then(transport.sendMessage(testMessage)))
|
||||
.expectError(McpTransportSessionNotFoundException.class)
|
||||
.verify(Duration.ofSeconds(5));
|
||||
|
||||
// Wait for second request to be made
|
||||
assertThat(secondRequestLatch.await(5, TimeUnit.SECONDS)).isTrue();
|
||||
|
||||
// Verify the second request included the session ID
|
||||
assertThat(lastReceivedSessionId.get()).isEqualTo("test-session-123");
|
||||
|
||||
// Verify exception handler was called with SessionNotFoundException using
|
||||
// timeout
|
||||
verify(exceptionHandler, timeout(5000)).accept(any(McpTransportSessionNotFoundException.class));
|
||||
}
|
||||
|
||||
/**
|
||||
* Test that 400 response WITHOUT session ID throws McpTransportException (not
|
||||
* SessionNotFoundException)
|
||||
*/
|
||||
@Test
|
||||
void test400WithoutSessionId() {
|
||||
serverResponseStatus.set(400);
|
||||
currentServerSessionId.set(null); // No session ID
|
||||
|
||||
var testMessage = createTestMessage();
|
||||
|
||||
StepVerifier.create(transport.sendMessage(testMessage))
|
||||
.expectErrorMatches(throwable -> throwable instanceof McpTransportException
|
||||
&& throwable.getMessage().contains("Bad Request") && throwable.getMessage().contains("400")
|
||||
&& !(throwable instanceof McpTransportSessionNotFoundException))
|
||||
.verify(Duration.ofSeconds(5));
|
||||
}
|
||||
|
||||
/**
|
||||
* Test that 400 response WITH session ID throws McpTransportSessionNotFoundException
|
||||
* Fixed version using proper async coordination
|
||||
*/
|
||||
@Test
|
||||
void test400WithSessionId() throws InterruptedException {
|
||||
|
||||
// First establish a session
|
||||
serverResponseStatus.set(200);
|
||||
currentServerSessionId.set("test-session-456");
|
||||
|
||||
// Set up exception handler
|
||||
@SuppressWarnings("unchecked")
|
||||
Consumer<Throwable> exceptionHandler = mock(Consumer.class);
|
||||
transport.setExceptionHandler(exceptionHandler);
|
||||
|
||||
// Connect with handler
|
||||
StepVerifier.create(transport.connect(msg -> msg)).verifyComplete();
|
||||
|
||||
// Send initial message to establish session
|
||||
var testMessage = createTestMessage();
|
||||
|
||||
// Send first message to establish session
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Wait for first request to complete
|
||||
boolean firstCompleted = firstRequestLatch.await(5, TimeUnit.SECONDS);
|
||||
assertThat(firstCompleted).isTrue();
|
||||
|
||||
// Wait for the GET request (SSE connection attempt) to complete
|
||||
boolean getCompleted = getRequestLatch.await(5, TimeUnit.SECONDS);
|
||||
assertThat(getCompleted).isTrue();
|
||||
|
||||
// Now return 400 for next request (simulating unknown session ID)
|
||||
serverResponseStatus.set(400);
|
||||
|
||||
// Use delaySubscription to ensure session is fully processed before next
|
||||
// request
|
||||
StepVerifier.create(Mono.delay(Duration.ofMillis(200)).then(transport.sendMessage(testMessage)))
|
||||
.expectError(McpTransportSessionNotFoundException.class)
|
||||
.verify(Duration.ofSeconds(5));
|
||||
|
||||
// Wait for second request to be made
|
||||
boolean secondCompleted = secondRequestLatch.await(5, TimeUnit.SECONDS);
|
||||
assertThat(secondCompleted).isTrue();
|
||||
|
||||
// Verify the second request included the session ID
|
||||
assertThat(lastReceivedSessionId.get()).isEqualTo("test-session-456");
|
||||
|
||||
// Verify exception handler was called with timeout
|
||||
verify(exceptionHandler, timeout(5000)).accept(any(McpTransportSessionNotFoundException.class));
|
||||
}
|
||||
|
||||
/**
|
||||
* Test session recovery after SessionNotFoundException Fixed version using reactive
|
||||
* patterns and proper synchronization
|
||||
*/
|
||||
@Test
|
||||
void testSessionRecoveryAfter404() {
|
||||
// First establish a session
|
||||
serverResponseStatus.set(200);
|
||||
currentServerSessionId.set("session-1");
|
||||
|
||||
// Send initial message to establish session
|
||||
var testMessage = createTestMessage();
|
||||
|
||||
// Use Mono.defer to ensure proper sequencing
|
||||
Mono<Void> establishSession = transport.sendMessage(testMessage).then(Mono.defer(() -> {
|
||||
// Simulate session loss - return 404
|
||||
serverResponseStatus.set(404);
|
||||
return transport.sendMessage(testMessage).onErrorResume(McpTransportSessionNotFoundException.class, e -> {
|
||||
// Expected error, continue with recovery
|
||||
return Mono.empty();
|
||||
});
|
||||
})).then(Mono.defer(() -> {
|
||||
// Now server is back with new session
|
||||
serverResponseStatus.set(200);
|
||||
currentServerSessionId.set("session-2");
|
||||
lastReceivedSessionId.set(null); // Reset to verify new session
|
||||
|
||||
// Should be able to establish new session
|
||||
return transport.sendMessage(testMessage);
|
||||
})).then(Mono.defer(() -> {
|
||||
// Verify no session ID was sent (since old session was invalidated)
|
||||
assertThat(lastReceivedSessionId.get()).isNull();
|
||||
|
||||
// Next request should use the new session ID
|
||||
return transport.sendMessage(testMessage);
|
||||
})).doOnSuccess(v -> {
|
||||
// Session ID should now be sent with requests
|
||||
assertThat(lastReceivedSessionId.get()).isEqualTo("session-2");
|
||||
});
|
||||
|
||||
StepVerifier.create(establishSession).verifyComplete();
|
||||
}
|
||||
|
||||
/**
|
||||
* Test that reconnect (GET request) also properly handles 404/400 errors Fixed
|
||||
* version with proper async handling
|
||||
*/
|
||||
@Test
|
||||
void testReconnectErrorHandling() throws InterruptedException {
|
||||
// Initialize latch for SSE connection
|
||||
CountDownLatch sseConnectionLatch = new CountDownLatch(1);
|
||||
|
||||
// Set up SSE endpoint for GET requests
|
||||
server.createContext("/mcp-sse", exchange -> {
|
||||
String method = exchange.getRequestMethod();
|
||||
String requestSessionId = exchange.getRequestHeaders().getFirst(HttpHeaders.MCP_SESSION_ID);
|
||||
|
||||
if ("GET".equals(method)) {
|
||||
sseConnectionLatch.countDown();
|
||||
int status = serverResponseStatus.get();
|
||||
|
||||
if (status == 404 && requestSessionId != null) {
|
||||
// 404 with session ID - should trigger SessionNotFoundException
|
||||
exchange.sendResponseHeaders(404, 0);
|
||||
}
|
||||
else if (status == 404) {
|
||||
// 404 without session ID - should trigger McpTransportException
|
||||
exchange.sendResponseHeaders(404, 0);
|
||||
}
|
||||
else {
|
||||
// Normal SSE response
|
||||
exchange.getResponseHeaders().set("Content-Type", "text/event-stream");
|
||||
exchange.sendResponseHeaders(200, 0);
|
||||
// Send a test SSE event
|
||||
String sseData = "event: message\ndata: {\"jsonrpc\":\"2.0\",\"method\":\"test\",\"params\":{}}\n\n";
|
||||
exchange.getResponseBody().write(sseData.getBytes());
|
||||
}
|
||||
}
|
||||
else {
|
||||
// POST request handling
|
||||
exchange.getResponseHeaders().set("Content-Type", "application/json");
|
||||
String responseSessionId = currentServerSessionId.get();
|
||||
if (responseSessionId != null) {
|
||||
exchange.getResponseHeaders().set(HttpHeaders.MCP_SESSION_ID, responseSessionId);
|
||||
}
|
||||
String response = "{\"jsonrpc\":\"2.0\",\"result\":{},\"id\":\"test-id\"}";
|
||||
exchange.sendResponseHeaders(200, response.length());
|
||||
exchange.getResponseBody().write(response.getBytes());
|
||||
}
|
||||
exchange.close();
|
||||
});
|
||||
|
||||
// Test with session ID - should get SessionNotFoundException
|
||||
serverResponseStatus.set(200);
|
||||
currentServerSessionId.set("sse-session-1");
|
||||
|
||||
var transport = WebClientStreamableHttpTransport.builder(WebClient.builder().baseUrl(HOST))
|
||||
.endpoint("/mcp-sse")
|
||||
.openConnectionOnStartup(true) // This will trigger GET request on connect
|
||||
.build();
|
||||
|
||||
// First connect successfully
|
||||
StepVerifier.create(transport.connect(msg -> msg)).verifyComplete();
|
||||
|
||||
// Wait for SSE connection to be established
|
||||
boolean connected = sseConnectionLatch.await(5, TimeUnit.SECONDS);
|
||||
assertThat(connected).isTrue();
|
||||
|
||||
// Send message to establish session
|
||||
var testMessage = createTestMessage();
|
||||
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
|
||||
|
||||
// Clean up
|
||||
StepVerifier.create(transport.closeGracefully()).verifyComplete();
|
||||
}
|
||||
|
||||
private McpSchema.JSONRPCRequest createTestMessage() {
|
||||
var initializeRequest = new McpSchema.InitializeRequest(ProtocolVersions.MCP_2025_03_26,
|
||||
McpSchema.ClientCapabilities.builder().roots(true).build(),
|
||||
new McpSchema.Implementation("Test Client", "1.0.0"));
|
||||
return new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, McpSchema.METHOD_INITIALIZE, "test-id",
|
||||
initializeRequest);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -11,6 +11,7 @@ import org.apache.catalina.LifecycleException;
|
|||
import org.apache.catalina.LifecycleState;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Configuration;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
|
@ -29,6 +30,7 @@ import io.modelcontextprotocol.server.McpServer.SingleSessionSyncSpecification;
|
|||
import io.modelcontextprotocol.server.transport.WebMvcSseServerTransportProvider;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
@Timeout(15)
|
||||
class WebMvcSseIntegrationTests extends AbstractMcpClientServerIntegrationTests {
|
||||
|
||||
private static final int PORT = TestUtil.findAvailablePort();
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import org.apache.catalina.LifecycleException;
|
|||
import org.apache.catalina.LifecycleState;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Configuration;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
|
@ -29,6 +30,7 @@ import io.modelcontextprotocol.server.McpServer.StatelessSyncSpecification;
|
|||
import io.modelcontextprotocol.server.transport.WebMvcStatelessServerTransport;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
@Timeout(15)
|
||||
class WebMvcStatelessIntegrationTests extends AbstractStatelessIntegrationTests {
|
||||
|
||||
private static final int PORT = TestUtil.findAvailablePort();
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import org.apache.catalina.LifecycleException;
|
|||
import org.apache.catalina.LifecycleState;
|
||||
import org.junit.jupiter.api.AfterEach;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.context.annotation.Configuration;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
|
|
@ -29,6 +30,7 @@ import io.modelcontextprotocol.server.McpServer.SyncSpecification;
|
|||
import io.modelcontextprotocol.server.transport.WebMvcStreamableServerTransportProvider;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
@Timeout(15)
|
||||
class WebMvcStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests {
|
||||
|
||||
private static final int PORT = TestUtil.findAvailablePort();
|
||||
|
|
|
|||
Loading…
Reference in a new issue