diff --git a/mcp-bom/pom.xml b/mcp-bom/pom.xml index 7214dacda..85cb15ea0 100644 --- a/mcp-bom/pom.xml +++ b/mcp-bom/pom.xml @@ -7,7 +7,7 @@ io.modelcontextprotocol.sdk mcp-parent - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT mcp-bom diff --git a/mcp-spring/mcp-spring-webflux/pom.xml b/mcp-spring/mcp-spring-webflux/pom.xml index 9008e52da..059108808 100644 --- a/mcp-spring/mcp-spring-webflux/pom.xml +++ b/mcp-spring/mcp-spring-webflux/pom.xml @@ -6,7 +6,7 @@ io.modelcontextprotocol.sdk mcp-parent - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT ../../pom.xml mcp-spring-webflux @@ -25,13 +25,13 @@ io.modelcontextprotocol.sdk mcp - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT io.modelcontextprotocol.sdk mcp-test - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT test diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransport.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransport.java index d7f7f9bfb..853aed2bf 100644 --- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransport.java +++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransport.java @@ -1,3 +1,7 @@ +/* + * Copyright 2025-2025 the original author or authors. + */ + package io.modelcontextprotocol.client.transport; import java.io.IOException; @@ -27,9 +31,11 @@ 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; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.Utils; import reactor.core.Disposable; @@ -65,9 +71,11 @@ */ 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 = "2025-03-26"; + private static final String MCP_PROTOCOL_VERSION = ProtocolVersions.MCP_2025_03_26; private static final String DEFAULT_ENDPOINT = "/mcp"; @@ -107,8 +115,8 @@ private WebClientStreamableHttpTransport(ObjectMapper objectMapper, WebClient.Bu } @Override - public String protocolVersion() { - return MCP_PROTOCOL_VERSION; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05, ProtocolVersions.MCP_2025_03_26); } /** @@ -138,11 +146,8 @@ private DefaultMcpTransportSession createTransportSession() { Function> onClose = sessionId -> sessionId == null ? Mono.empty() : webClient.delete() .uri(this.endpoint) + .header(HttpHeaders.MCP_SESSION_ID, sessionId) .header(HttpHeaders.PROTOCOL_VERSION, MCP_PROTOCOL_VERSION) - .headers(httpHeaders -> { - httpHeaders.add(HttpHeaders.MCP_SESSION_ID, sessionId); - httpHeaders.add(HttpHeaders.PROTOCOL_VERSION, MCP_PROTOCOL_VERSION); - }) .retrieve() .toBodilessEntity() .onErrorComplete(e -> { @@ -219,8 +224,13 @@ else if (isNotAllowed(response)) { 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.createError().doOnError(e -> { @@ -316,10 +326,10 @@ else if (mediaType.isCompatibleWith(MediaType.APPLICATION_JSON)) { } } 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))) @@ -360,10 +370,10 @@ private Flux extractError(ClientResponse response, Str 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); } @@ -372,7 +382,11 @@ private Flux extractError(ClientResponse response, Str // 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(); @@ -401,7 +415,7 @@ private static boolean isEventStream(ClientResponse response) { } private static String sessionIdOrPlaceholder(McpTransportSession transportSession) { - return transportSession.sessionId().orElse("[missing_session_id]"); + return transportSession.sessionId().orElse(MISSING_SESSION_ID); } private Flux directResponseFlux(McpSchema.JSONRPCMessage sentMessage, @@ -419,8 +433,7 @@ private Flux directResponseFlux(McpSchema.JSONRPCMessa } } catch (IOException e) { - // TODO: this should be a McpTransportError - s.error(e); + s.error(new McpTransportException(e)); } }).flatMapIterable(Function.identity()); } @@ -447,7 +460,7 @@ private Tuple2, Iterable> parse(Serve 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 { diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransport.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransport.java index fe6b07a6d..51d21d18b 100644 --- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransport.java +++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/client/transport/WebFluxSseClientTransport.java @@ -1,9 +1,11 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol.client.transport; import java.io.IOException; +import java.util.List; import java.util.function.BiConsumer; import java.util.function.Function; @@ -12,9 +14,9 @@ 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; import io.modelcontextprotocol.util.Assert; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -64,7 +66,7 @@ public class WebFluxSseClientTransport implements McpClientTransport { private static final Logger logger = LoggerFactory.getLogger(WebFluxSseClientTransport.class); - private static final String MCP_PROTOCOL_VERSION = "2024-11-05"; + private static final String MCP_PROTOCOL_VERSION = ProtocolVersions.MCP_2024_11_05; /** * Event type for JSON-RPC messages received through the SSE connection. The server @@ -171,8 +173,8 @@ public WebFluxSseClientTransport(WebClient.Builder webClientBuilder, ObjectMappe } @Override - public String protocolVersion() { - return MCP_PROTOCOL_VERSION; + public List protocolVersions() { + return List.of(MCP_PROTOCOL_VERSION); } /** @@ -194,8 +196,6 @@ public String protocolVersion() { * @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 connect(Function, Mono> handler) { @@ -212,7 +212,7 @@ public Mono connect(Function, Mono> h 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())) { diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java index 67810fb56..ead7380f0 100644 --- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java +++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxSseServerTransportProvider.java @@ -1,16 +1,26 @@ +/* + * Copyright 2025-2025 the original author or authors. + */ + package io.modelcontextprotocol.server.transport; import java.io.IOException; import java.time.Duration; +import java.util.List; import java.util.concurrent.ConcurrentHashMap; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; + +import io.modelcontextprotocol.server.DefaultMcpTransportContext; +import io.modelcontextprotocol.server.McpTransportContext; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpServerSession; import io.modelcontextprotocol.spec.McpServerTransport; import io.modelcontextprotocol.spec.McpServerTransportProvider; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.KeepAliveScheduler; @@ -109,6 +119,8 @@ public class WebFluxSseServerTransportProvider implements McpServerTransportProv */ private final ConcurrentHashMap sessions = new ConcurrentHashMap<>(); + private McpTransportContextExtractor contextExtractor; + /** * Flag indicating if the transport is shutting down. */ @@ -188,15 +200,38 @@ public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String baseU @Deprecated public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, String sseEndpoint, Duration keepAliveInterval) { + this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, keepAliveInterval, + (serverRequest, context) -> context); + } + + /** + * Constructs a new WebFlux SSE server transport provider instance. + * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization + * of MCP messages. Must not be null. + * @param baseUrl webflux message base path + * @param messageEndpoint The endpoint URI where clients should send their JSON-RPC + * messages. This endpoint will be communicated to clients during SSE connection + * setup. Must not be null. + * @param sseEndpoint The SSE endpoint path. Must not be null. + * @param keepAliveInterval The interval for sending keep-alive pings to clients. + * @param contextExtractor The context extractor to use for extracting MCP transport + * context from HTTP requests. Must not be null. + * @throws IllegalArgumentException if either parameter is null + */ + private WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, + String sseEndpoint, Duration keepAliveInterval, + McpTransportContextExtractor contextExtractor) { Assert.notNull(objectMapper, "ObjectMapper must not be null"); Assert.notNull(baseUrl, "Message base path must not be null"); Assert.notNull(messageEndpoint, "Message endpoint must not be null"); Assert.notNull(sseEndpoint, "SSE endpoint must not be null"); + Assert.notNull(contextExtractor, "Context extractor must not be null"); this.objectMapper = objectMapper; this.baseUrl = baseUrl; this.messageEndpoint = messageEndpoint; this.sseEndpoint = sseEndpoint; + this.contextExtractor = contextExtractor; this.routerFunction = RouterFunctions.route() .GET(this.sseEndpoint, this::handleSseConnection) .POST(this.messageEndpoint, this::handleMessage) @@ -215,8 +250,8 @@ public WebFluxSseServerTransportProvider(ObjectMapper objectMapper, String baseU } @Override - public String protocolVersion() { - return "2024-11-05"; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05); } @Override @@ -309,6 +344,8 @@ private Mono handleSseConnection(ServerRequest request) { return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).bodyValue("Server is shutting down"); } + McpTransportContext transportContext = this.contextExtractor.extract(request, new DefaultMcpTransportContext()); + return ServerResponse.ok() .contentType(MediaType.TEXT_EVENT_STREAM) .body(Flux.>create(sink -> { @@ -330,7 +367,7 @@ private Mono handleSseConnection(ServerRequest request) { logger.debug("Session {} cancelled", sessionId); sessions.remove(sessionId); }); - }), ServerSentEvent.class); + }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)), ServerSentEvent.class); } /** @@ -364,6 +401,8 @@ private Mono handleMessage(ServerRequest request) { .bodyValue(new McpError("Session not found: " + request.queryParam("sessionId").get())); } + McpTransportContext transportContext = this.contextExtractor.extract(request, new DefaultMcpTransportContext()); + return request.bodyToMono(String.class).flatMap(body -> { try { McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body); @@ -380,7 +419,7 @@ private Mono handleMessage(ServerRequest request) { logger.error("Failed to deserialize message: {}", e.getMessage()); return ServerResponse.badRequest().bodyValue(new McpError("Invalid message format")); } - }); + }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)); } private class WebFluxMcpSessionTransport implements McpServerTransport { @@ -452,6 +491,8 @@ public static class Builder { private Duration keepAliveInterval; + private McpTransportContextExtractor contextExtractor = (serverRequest, context) -> context; + /** * Sets the ObjectMapper to use for JSON serialization/deserialization of MCP * messages. @@ -513,6 +554,22 @@ public Builder keepAliveInterval(Duration keepAliveInterval) { return this; } + /** + * Sets the context extractor that allows providing the MCP feature + * implementations to inspect HTTP transport level metadata that was present at + * HTTP request processing time. This allows to extract custom headers and other + * useful data for use during execution later on in the process. + * @param contextExtractor The contextExtractor to fill in a + * {@link McpTransportContext}. + * @return this builder instance + * @throws IllegalArgumentException if contextExtractor is null + */ + public Builder contextExtractor(McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "contextExtractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + /** * Builds a new instance of {@link WebFluxSseServerTransportProvider} with the * configured settings. @@ -524,7 +581,7 @@ public WebFluxSseServerTransportProvider build() { Assert.notNull(messageEndpoint, "Message endpoint must be set"); return new WebFluxSseServerTransportProvider(objectMapper, baseUrl, messageEndpoint, sseEndpoint, - keepAliveInterval); + keepAliveInterval, contextExtractor); } } diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStatelessServerTransport.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStatelessServerTransport.java index c514f2dff..23fff25b3 100644 --- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStatelessServerTransport.java +++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStatelessServerTransport.java @@ -1,3 +1,7 @@ +/* + * Copyright 2025-2025 the original author or authors. + */ + package io.modelcontextprotocol.server.transport; import com.fasterxml.jackson.databind.ObjectMapper; diff --git a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java index 00ec68c5d..963a50249 100644 --- a/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java +++ b/mcp-spring/mcp-spring-webflux/src/main/java/io/modelcontextprotocol/server/transport/WebFluxStreamableServerTransportProvider.java @@ -1,3 +1,7 @@ +/* + * Copyright 2025-2025 the original author or authors. + */ + package io.modelcontextprotocol.server.transport; import com.fasterxml.jackson.core.type.TypeReference; @@ -10,6 +14,7 @@ import io.modelcontextprotocol.spec.McpStreamableServerSession; import io.modelcontextprotocol.spec.McpStreamableServerTransport; import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.server.McpTransportContext; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.KeepAliveScheduler; @@ -89,15 +94,11 @@ private WebFluxStreamableServerTransportProvider(ObjectMapper objectMapper, Stri this.keepAliveScheduler.start(); } - else { - logger.warn("Keep-alive interval is not set or invalid. No keep-alive will be scheduled."); - } - } @Override - public String protocolVersion() { - return "2025-03-26"; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05, ProtocolVersions.MCP_2025_03_26); } @Override @@ -190,7 +191,9 @@ private Mono handleGet(ServerRequest request) { String lastId = request.headers().asHttpHeaders().getFirst(HttpHeaders.LAST_EVENT_ID); return ServerResponse.ok() .contentType(MediaType.TEXT_EVENT_STREAM) - .body(session.replay(lastId), ServerSentEvent.class); + .body(session.replay(lastId) + .contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)), + ServerSentEvent.class); } return ServerResponse.ok() @@ -201,7 +204,9 @@ private Mono handleGet(ServerRequest request) { McpStreamableServerSession.McpStreamableServerSessionStream listeningStream = session .listeningStream(sessionTransport); sink.onDispose(listeningStream::close); - }), ServerSentEvent.class); + // TODO Clarify why the outer context is not present in the + // Flux.create sink? + }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)), ServerSentEvent.class); }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)); } @@ -281,7 +286,10 @@ else if (message instanceof McpSchema.JSONRPCRequest jsonrpcRequest) { return true; }).contextWrite(sink.contextView()).subscribe(); sink.onCancel(streamSubscription); - }), ServerSentEvent.class); + // TODO Clarify why the outer context is not present in the + // Flux.create sink? + }).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)), + ServerSentEvent.class); } else { return ServerResponse.badRequest().bodyValue(new McpError("Unknown message type")); diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java index 23ddf6173..c8dc6e90b 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxSseIntegrationTests.java @@ -1,37 +1,19 @@ /* * Copyright 2024 - 2024 the original author or authors. */ -package io.modelcontextprotocol; -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatExceptionOfType; -import static org.assertj.core.api.Assertions.assertWith; -import static org.awaitility.Awaitility.await; -import static org.mockito.Mockito.mock; +package io.modelcontextprotocol; import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.BiFunction; -import java.util.function.Function; -import java.util.stream.Collectors; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; +import org.junit.jupiter.api.Timeout; import org.springframework.http.server.reactive.HttpHandler; import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; -import org.springframework.web.client.RestClient; import org.springframework.web.reactive.function.client.WebClient; import org.springframework.web.reactive.function.server.RouterFunctions; +import org.springframework.web.reactive.function.server.ServerRequest; import com.fasterxml.jackson.databind.ObjectMapper; @@ -39,36 +21,16 @@ import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport; import io.modelcontextprotocol.client.transport.WebFluxSseClientTransport; import io.modelcontextprotocol.server.McpServer; -import io.modelcontextprotocol.server.McpServerFeatures; -import io.modelcontextprotocol.server.McpSyncServerExchange; +import io.modelcontextprotocol.server.McpServer.AsyncSpecification; +import io.modelcontextprotocol.server.McpServer.SingleSessionSyncSpecification; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.server.TestUtil; import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider; -import io.modelcontextprotocol.spec.McpError; -import io.modelcontextprotocol.spec.McpSchema; -import io.modelcontextprotocol.spec.McpSchema.CallToolResult; -import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities; -import io.modelcontextprotocol.spec.McpSchema.CompleteRequest; -import io.modelcontextprotocol.spec.McpSchema.CompleteResult; -import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest; -import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult; -import io.modelcontextprotocol.spec.McpSchema.ElicitRequest; -import io.modelcontextprotocol.spec.McpSchema.ElicitResult; -import io.modelcontextprotocol.spec.McpSchema.InitializeResult; -import io.modelcontextprotocol.spec.McpSchema.ModelPreferences; -import io.modelcontextprotocol.spec.McpSchema.Prompt; -import io.modelcontextprotocol.spec.McpSchema.PromptArgument; -import io.modelcontextprotocol.spec.McpSchema.PromptReference; -import io.modelcontextprotocol.spec.McpSchema.Role; -import io.modelcontextprotocol.spec.McpSchema.Root; -import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; -import io.modelcontextprotocol.spec.McpSchema.Tool; -import net.javacrumbs.jsonunit.core.Option; -import reactor.core.publisher.Mono; import reactor.netty.DisposableServer; import reactor.netty.http.server.HttpServer; -import reactor.test.StepVerifier; -class WebFluxSseIntegrationTests { +@Timeout(15) +class WebFluxSseIntegrationTests extends AbstractMcpClientServerIntegrationTests { private static final int PORT = TestUtil.findAvailablePort(); @@ -80,1413 +42,61 @@ class WebFluxSseIntegrationTests { private WebFluxSseServerTransportProvider mcpServerTransportProvider; - ConcurrentHashMap clientBuilders = new ConcurrentHashMap<>(); + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; - @BeforeEach - public void before() { + @Override + protected void prepareClients(int port, String mcpEndpoint) { - this.mcpServerTransportProvider = new WebFluxSseServerTransportProvider.Builder() - .objectMapper(new ObjectMapper()) - .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) - .sseEndpoint(CUSTOM_SSE_ENDPOINT) - .build(); - - HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpServerTransportProvider.getRouterFunction()); - ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler); - this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow(); + clientBuilders + .put("httpclient", + McpClient.sync(HttpClientSseClientTransport.builder("http://localhost:" + PORT) + .sseEndpoint(CUSTOM_SSE_ENDPOINT) + .build()).requestTimeout(Duration.ofHours(10))); - clientBuilders.put("httpclient", - McpClient.sync(HttpClientSseClientTransport.builder("http://localhost:" + PORT) - .sseEndpoint(CUSTOM_SSE_ENDPOINT) - .build())); clientBuilders.put("webflux", McpClient .sync(WebFluxSseClientTransport.builder(WebClient.builder().baseUrl("http://localhost:" + PORT)) .sseEndpoint(CUSTOM_SSE_ENDPOINT) - .build())); - - } - - @AfterEach - public void after() { - if (httpServer != null) { - httpServer.disposeNow(); - } - } - - // --------------------------------------- - // Sampling Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateMessageWithoutSamplingCapabilities(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> exchange.createMessage(mock(CreateMessageRequest.class)) - .thenReturn(mock(CallToolResult.class))) - .build(); - - var server = McpServer.async(mcpServerTransportProvider).serverInfo("test-server", "1.0.0").tools(tool).build(); - - try (var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample " + "client", "0.0.0")) - .build();) { - - assertThat(client.initialize()).isNotNull(); - - try { - client.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - } - catch (McpError e) { - assertThat(e).isInstanceOf(McpError.class) - .hasMessage("Client must be configured with sampling capabilities"); - } - } - server.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateMessageSuccess(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - Function samplingHandler = request -> { - assertThat(request.messages()).hasSize(1); - assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); - - return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", - CreateMessageResult.StopReason.STOP_SEQUENCE); - }; - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - AtomicReference samplingResult = new AtomicReference<>(); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var createMessageRequest = McpSchema.CreateMessageRequest.builder() - .messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, - new McpSchema.TextContent("Test message")))) - .modelPreferences(ModelPreferences.builder() - .hints(List.of()) - .costPriority(1.0) - .speedPriority(1.0) - .intelligencePriority(1.0) .build()) - .build(); - - return exchange.createMessage(createMessageRequest) - .doOnNext(samplingResult::set) - .thenReturn(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - - assertWith(samplingResult.get(), result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }); - } - mcpServer.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws InterruptedException { - - // Client - var clientBuilder = clientBuilders.get(clientType); - - Function samplingHandler = request -> { - assertThat(request.messages()).hasSize(1); - assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", - CreateMessageResult.StopReason.STOP_SEQUENCE); - }; - - // Server - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - AtomicReference samplingResult = new AtomicReference<>(); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var craeteMessageRequest = McpSchema.CreateMessageRequest.builder() - .messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, - new McpSchema.TextContent("Test message")))) - .modelPreferences(ModelPreferences.builder() - .hints(List.of()) - .costPriority(1.0) - .speedPriority(1.0) - .intelligencePriority(1.0) - .build()) - .build(); - - return exchange.createMessage(craeteMessageRequest) - .doOnNext(samplingResult::set) - .thenReturn(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .requestTimeout(Duration.ofSeconds(4)) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - - assertWith(samplingResult.get(), result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }); - } - - mcpServer.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateMessageWithRequestTimeoutFail(String clientType) throws InterruptedException { - - // Client - var clientBuilder = clientBuilders.get(clientType); - - Function samplingHandler = request -> { - assertThat(request.messages()).hasSize(1); - assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", - CreateMessageResult.StopReason.STOP_SEQUENCE); - }; - - // Server - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var craeteMessageRequest = McpSchema.CreateMessageRequest.builder() - .messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, - new McpSchema.TextContent("Test message")))) - .build(); - - return exchange.createMessage(craeteMessageRequest).thenReturn(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .requestTimeout(Duration.ofSeconds(1)) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThatExceptionOfType(McpError.class).isThrownBy(() -> { - mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - }).withMessageContaining("within 1000ms"); - - } - - mcpServer.closeGracefully().block(); - } - - // --------------------------------------- - // Elicitation Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateElicitationWithoutElicitationCapabilities(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - exchange.createElicitation(mock(ElicitRequest.class)).block(); - - return Mono.just(mock(CallToolResult.class)); - }) - .build(); - - var server = McpServer.async(mcpServerTransportProvider).serverInfo("test-server", "1.0.0").tools(tool).build(); - - try ( - // Create client without elicitation capabilities - var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) { - - assertThat(client.initialize()).isNotNull(); - - try { - client.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - } - catch (McpError e) { - assertThat(e).isInstanceOf(McpError.class) - .hasMessage("Client must be configured with elicitation capabilities"); - } - } - server.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateElicitationSuccess(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - Function elicitationHandler = request -> { - assertThat(request.message()).isNotEmpty(); - assertThat(request.requestedSchema()).isNotNull(); - - return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); - }; - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var elicitationRequest = ElicitRequest.builder() - .message("Test message") - .requestedSchema( - Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) - .build(); - - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().elicitation().build()) - .elicitation(elicitationHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - } - mcpServer.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { - - // Client - var clientBuilder = clientBuilders.get(clientType); - - Function elicitationHandler = request -> { - assertThat(request.message()).isNotEmpty(); - assertThat(request.requestedSchema()).isNotNull(); - return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); - }; - - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().elicitation().build()) - .elicitation(elicitationHandler) - .build(); - - // Server - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var elicitationRequest = ElicitRequest.builder() - .message("Test message") - .requestedSchema( - Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) - .build(); - - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(3)) - .tools(tool) - .build(); - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - - mcpClient.closeGracefully(); - mcpServer.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateElicitationWithRequestTimeoutFail(String clientType) { - - var latch = new CountDownLatch(1); - // Client - var clientBuilder = clientBuilders.get(clientType); - - Function elicitationHandler = request -> { - assertThat(request.message()).isNotEmpty(); - assertThat(request.requestedSchema()).isNotNull(); - - try { - if (!latch.await(2, TimeUnit.SECONDS)) { - throw new RuntimeException("Timeout waiting for elicitation processing"); - } - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); - }; - - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().elicitation().build()) - .elicitation(elicitationHandler) - .build(); - - // Server - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var elicitationRequest = ElicitRequest.builder() - .message("Test message") - .requestedSchema( - Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) - .build(); - - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(1)) // 1 second. - .tools(tool) - .build(); - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThatExceptionOfType(McpError.class).isThrownBy(() -> { - mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - }).withMessageContaining("within 1000ms"); - - mcpClient.closeGracefully(); - mcpServer.closeGracefully().block(); - } - - // --------------------------------------- - // Roots Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsSuccess(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - List roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2")); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(roots) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(rootsRef.get()).isNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(roots); - }); - - // Remove a root - mcpClient.removeRoot(roots.get(0).uri()); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(roots.get(1))); - }); - - // Add a new root - var root3 = new Root("uri3://", "root3"); - mcpClient.addRoot(root3); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(roots.get(1), root3)); - }); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsWithoutCapability(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - McpServerFeatures.SyncToolSpecification tool = McpServerFeatures.SyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - exchange.listRoots(); // try to list roots - - return mock(CallToolResult.class); - }) - .build(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider).rootsChangeHandler((exchange, rootsUpdate) -> { - }).tools(tool).build(); - - // Create client without roots capability - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()).build()) { - - assertThat(mcpClient.initialize()).isNotNull(); - - // Attempt to list roots should fail - try { - mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - } - catch (McpError e) { - assertThat(e).isInstanceOf(McpError.class).hasMessage("Roots not supported"); - } - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsNotificationWithEmptyRootsList(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(List.of()) // Empty roots list - .build()) { - - assertThat(mcpClient.initialize()).isNotNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).isEmpty(); - }); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsWithMultipleHandlers(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); + .requestTimeout(Duration.ofHours(10))); - List roots = List.of(new Root("uri1://", "root1")); - - AtomicReference> rootsRef1 = new AtomicReference<>(); - AtomicReference> rootsRef2 = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef1.set(rootsUpdate)) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef2.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(roots) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef1.get()).containsAll(roots); - assertThat(rootsRef2.get()).containsAll(roots); - }); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsServerCloseWithActiveSubscription(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - List roots = List.of(new Root("uri1://", "root1")); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(roots) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(roots); - }); - } - - mcpServer.close(); - } - - // --------------------------------------- - // Tools Tests - // --------------------------------------- - - String emptyJsonSchema = """ - { - "$schema": "http://json-schema.org/draft-07/schema#", - "type": "object", - "properties": {} - } - """; - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testToolCallSuccess(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - return callResponse; - }) - .build(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool1) - .build(); - - try (var mcpClient = clientBuilder.build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testToolListChangeHandlingSuccess(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - return callResponse; - }) - .build(); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool1) - .build(); - - try (var mcpClient = clientBuilder.toolsChangeConsumer(toolsUpdate -> { - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - rootsRef.set(toolsUpdate); - }).build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(rootsRef.get()).isNull(); - - assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); - - mcpServer.notifyToolsListChanged(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool1.tool())); - }); - - // Remove a tool - mcpServer.removeTool("tool1"); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).isEmpty(); - }); - - // Add a new tool - McpServerFeatures.SyncToolSpecification tool2 = McpServerFeatures.SyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool2", "tool2 description", emptyJsonSchema)) - .callHandler((exchange, request) -> callResponse) - .build(); - - mcpServer.addTool(tool2); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool2.tool())); - }); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testInitialize(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var mcpServer = McpServer.sync(mcpServerTransportProvider).build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - } - - mcpServer.close(); - } - - // --------------------------------------- - // Logging Tests - // --------------------------------------- - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testLoggingNotification(String clientType) throws InterruptedException { - int expectedNotificationsCount = 3; - CountDownLatch latch = new CountDownLatch(expectedNotificationsCount); - // Create a list to store received logging notifications - List receivedNotifications = new CopyOnWriteArrayList<>(); - - var clientBuilder = clientBuilders.get(clientType); - - // Create server with a tool that sends logging notifications - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("logging-test", "Test logging notifications", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - // Create and send notifications with different levels - - //@formatter:off - return exchange // This should be filtered out (DEBUG < NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.DEBUG) - .logger("test-logger") - .data("Debug message") - .build()) - .then(exchange // This should be sent (NOTICE >= NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.NOTICE) - .logger("test-logger") - .data("Notice message") - .build())) - .then(exchange // This should be sent (ERROR > NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.ERROR) - .logger("test-logger") - .data("Error message") - .build())) - .then(exchange // This should be filtered out (INFO < NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.INFO) - .logger("test-logger") - .data("Another info message") - .build())) - .then(exchange // This should be sent (ERROR >= NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.ERROR) - .logger("test-logger") - .data("Another error message") - .build())) - .thenReturn(new CallToolResult("Logging test completed", false)); - //@formatter:on - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().logging().tools(true).build()) - .tools(tool) - .build(); - - try ( - // Create client with logging notification handler - var mcpClient = clientBuilder.loggingConsumer(notification -> { - receivedNotifications.add(notification); - latch.countDown(); - }).build()) { - - // Initialize client - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Set minimum logging level to NOTICE - mcpClient.setLoggingLevel(McpSchema.LoggingLevel.NOTICE); - - // Call the tool that sends logging notifications - CallToolResult result = mcpClient.callTool(new McpSchema.CallToolRequest("logging-test", Map.of())); - assertThat(result).isNotNull(); - assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Logging test completed"); - - assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue(); - - // Should have received 3 notifications (1 NOTICE and 2 ERROR) - assertThat(receivedNotifications).hasSize(expectedNotificationsCount); - - Map notificationMap = receivedNotifications.stream() - .collect(Collectors.toMap(n -> n.data(), n -> n)); - - // First notification should be NOTICE level - assertThat(notificationMap.get("Notice message").level()).isEqualTo(McpSchema.LoggingLevel.NOTICE); - assertThat(notificationMap.get("Notice message").logger()).isEqualTo("test-logger"); - assertThat(notificationMap.get("Notice message").data()).isEqualTo("Notice message"); - - // Second notification should be ERROR level - assertThat(notificationMap.get("Error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); - assertThat(notificationMap.get("Error message").logger()).isEqualTo("test-logger"); - assertThat(notificationMap.get("Error message").data()).isEqualTo("Error message"); - - // Third notification should be ERROR level - assertThat(notificationMap.get("Another error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); - assertThat(notificationMap.get("Another error message").logger()).isEqualTo("test-logger"); - assertThat(notificationMap.get("Another error message").data()).isEqualTo("Another error message"); - } - mcpServer.close(); - } - - // --------------------------------------- - // Progress Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testProgressNotification(String clientType) throws InterruptedException { - int expectedNotificationsCount = 4; // 3 notifications + 1 for another progress - // token - CountDownLatch latch = new CountDownLatch(expectedNotificationsCount); - // Create a list to store received logging notifications - List receivedNotifications = new CopyOnWriteArrayList<>(); - - var clientBuilder = clientBuilders.get(clientType); - - // Create server with a tool that sends logging notifications - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(McpSchema.Tool.builder() - .name("progress-test") - .description("Test progress notifications") - .inputSchema(emptyJsonSchema) - .build()) - .callHandler((exchange, request) -> { - - // Create and send notifications - var progressToken = (String) request.meta().get("progressToken"); - - return exchange - .progressNotification( - new McpSchema.ProgressNotification(progressToken, 0.0, 1.0, "Processing started")) - .then(exchange.progressNotification( - new McpSchema.ProgressNotification(progressToken, 0.5, 1.0, "Processing data"))) - .then(// Send a progress notification with another progress value - // should - exchange.progressNotification(new McpSchema.ProgressNotification("another-progress-token", - 0.0, 1.0, "Another processing started"))) - .then(exchange.progressNotification( - new McpSchema.ProgressNotification(progressToken, 1.0, 1.0, "Processing completed"))) - .thenReturn(new CallToolResult(("Progress test completed"), false)); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try ( - // Create client with progress notification handler - var mcpClient = clientBuilder.progressConsumer(notification -> { - receivedNotifications.add(notification); - latch.countDown(); - }).build()) { - - // Initialize client - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call the tool that sends progress notifications - McpSchema.CallToolRequest callToolRequest = McpSchema.CallToolRequest.builder() - .name("progress-test") - .meta(Map.of("progressToken", "test-progress-token")) - .build(); - CallToolResult result = mcpClient.callTool(callToolRequest); - assertThat(result).isNotNull(); - assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Progress test completed"); - - assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue(); - - // Should have received 3 notifications - assertThat(receivedNotifications).hasSize(expectedNotificationsCount); - - Map notificationMap = receivedNotifications.stream() - .collect(Collectors.toMap(n -> n.message(), n -> n)); - - // First notification should be 0.0/1.0 progress - assertThat(notificationMap.get("Processing started").progressToken()).isEqualTo("test-progress-token"); - assertThat(notificationMap.get("Processing started").progress()).isEqualTo(0.0); - assertThat(notificationMap.get("Processing started").total()).isEqualTo(1.0); - assertThat(notificationMap.get("Processing started").message()).isEqualTo("Processing started"); - - // Second notification should be 0.5/1.0 progress - assertThat(notificationMap.get("Processing data").progressToken()).isEqualTo("test-progress-token"); - assertThat(notificationMap.get("Processing data").progress()).isEqualTo(0.5); - assertThat(notificationMap.get("Processing data").total()).isEqualTo(1.0); - assertThat(notificationMap.get("Processing data").message()).isEqualTo("Processing data"); - - // Third notification should be another progress token with 0.0/1.0 progress - assertThat(notificationMap.get("Another processing started").progressToken()) - .isEqualTo("another-progress-token"); - assertThat(notificationMap.get("Another processing started").progress()).isEqualTo(0.0); - assertThat(notificationMap.get("Another processing started").total()).isEqualTo(1.0); - assertThat(notificationMap.get("Another processing started").message()) - .isEqualTo("Another processing started"); - - // Fourth notification should be 1.0/1.0 progress - assertThat(notificationMap.get("Processing completed").progressToken()).isEqualTo("test-progress-token"); - assertThat(notificationMap.get("Processing completed").progress()).isEqualTo(1.0); - assertThat(notificationMap.get("Processing completed").total()).isEqualTo(1.0); - assertThat(notificationMap.get("Processing completed").message()).isEqualTo("Processing completed"); - } - finally { - mcpServer.close(); - } } - // --------------------------------------- - // Completion Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : Completion call") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCompletionShouldReturnExpectedSuggestions(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - var expectedValues = List.of("python", "pytorch", "pyside"); - var completionResponse = new McpSchema.CompleteResult(new CompleteResult.CompleteCompletion(expectedValues, 10, // total - true // hasMore - )); - - AtomicReference samplingRequest = new AtomicReference<>(); - BiFunction completionHandler = (mcpSyncServerExchange, - request) -> { - samplingRequest.set(request); - return completionResponse; - }; - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .capabilities(ServerCapabilities.builder().completions().build()) - .prompts(new McpServerFeatures.SyncPromptSpecification( - new Prompt("code_review", "Code review", "this is code review prompt", - List.of(new PromptArgument("language", "Language", "string", false))), - (mcpSyncServerExchange, getPromptRequest) -> null)) - .completions(new McpServerFeatures.SyncCompletionSpecification( - new McpSchema.PromptReference("ref/prompt", "code_review", "Code review"), completionHandler)) - .build(); - - try (var mcpClient = clientBuilder.build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CompleteRequest request = new CompleteRequest( - new PromptReference("ref/prompt", "code_review", "Code review"), - new CompleteRequest.CompleteArgument("language", "py")); - - CompleteResult result = mcpClient.completeCompletion(request); - - assertThat(result).isNotNull(); - - assertThat(samplingRequest.get().argument().name()).isEqualTo("language"); - assertThat(samplingRequest.get().argument().value()).isEqualTo("py"); - assertThat(samplingRequest.get().ref().type()).isEqualTo("ref/prompt"); - } - - mcpServer.close(); + @Override + protected AsyncSpecification prepareAsyncServerBuilder() { + return McpServer.async(mcpServerTransportProvider); } - // --------------------------------------- - // Ping Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testPingSuccess(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create server with a tool that uses ping functionality - AtomicReference executionOrder = new AtomicReference<>(""); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("ping-async-test", "Test ping async behavior", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - executionOrder.set(executionOrder.get() + "1"); - - // Test async ping behavior - return exchange.ping().doOnNext(result -> { - - assertThat(result).isNotNull(); - // Ping should return an empty object or map - assertThat(result).isInstanceOf(Map.class); - - executionOrder.set(executionOrder.get() + "2"); - assertThat(result).isNotNull(); - }).then(Mono.fromCallable(() -> { - executionOrder.set(executionOrder.get() + "3"); - return new CallToolResult("Async ping test completed", false); - })); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - - // Initialize client - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call the tool that tests ping async behavior - CallToolResult result = mcpClient.callTool(new McpSchema.CallToolRequest("ping-async-test", Map.of())); - assertThat(result).isNotNull(); - assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Async ping test completed"); - - // Verify execution order - assertThat(executionOrder.get()).isEqualTo("123"); - } - - mcpServer.closeGracefully().block(); + @Override + protected SingleSessionSyncSpecification prepareSyncServerBuilder() { + return McpServer.sync(mcpServerTransportProvider); } - // --------------------------------------- - // Tool Structured Output Schema Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputValidationSuccess(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create a tool with output schema - Map outputSchema = Map.of( - "type", "object", "properties", Map.of("result", Map.of("type", "number"), "operation", - Map.of("type", "string"), "timestamp", Map.of("type", "string")), - "required", List.of("result", "operation")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification tool = new McpServerFeatures.SyncToolSpecification(calculatorTool, - (exchange, request) -> { - String expression = (String) request.getOrDefault("expression", "2 + 3"); - double result = evaluateExpression(expression); - return CallToolResult.builder() - .structuredContent( - Map.of("result", result, "operation", expression, "timestamp", "2024-01-01T10:00:00Z")) - .build(); - }); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Verify tool is listed with output schema - var toolsList = mcpClient.listTools(); - assertThat(toolsList.tools()).hasSize(1); - assertThat(toolsList.tools().get(0).name()).isEqualTo("calculator"); - // Note: outputSchema might be null in sync server, but validation still works - - // Call tool with valid structured output - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isFalse(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - assertThatJson(((McpSchema.TextContent) response.content().get(0)).text()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"result":5.0,"operation":"2 + 3","timestamp":"2024-01-01T10:00:00Z"}""")); - - assertThat(response.structuredContent()).isNotNull(); - assertThatJson(response.structuredContent()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"result":5.0,"operation":"2 + 3","timestamp":"2024-01-01T10:00:00Z"}""")); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputValidationFailure(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create a tool with output schema - Map outputSchema = Map.of("type", "object", "properties", - Map.of("result", Map.of("type", "number"), "operation", Map.of("type", "string")), "required", - List.of("result", "operation")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification tool = new McpServerFeatures.SyncToolSpecification(calculatorTool, - (exchange, request) -> { - // Return invalid structured output. Result should be number, missing - // operation - return CallToolResult.builder() - .addTextContent("Invalid calculation") - .structuredContent(Map.of("result", "not-a-number", "extra", "field")) - .build(); - }); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call tool with invalid structured output - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isTrue(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - String errorMessage = ((McpSchema.TextContent) response.content().get(0)).text(); - assertThat(errorMessage).contains("Validation failed"); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputMissingStructuredContent(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create a tool with output schema - Map outputSchema = Map.of("type", "object", "properties", - Map.of("result", Map.of("type", "number")), "required", List.of("result")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification tool = new McpServerFeatures.SyncToolSpecification(calculatorTool, - (exchange, request) -> { - // Return result without structured content but tool has output schema - return CallToolResult.builder().addTextContent("Calculation completed").build(); - }); + @BeforeEach + public void before() { - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) + this.mcpServerTransportProvider = new WebFluxSseServerTransportProvider.Builder() + .objectMapper(new ObjectMapper()) + .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) + .sseEndpoint(CUSTOM_SSE_ENDPOINT) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .build(); - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call tool that should return structured content but doesn't - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isTrue(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - String errorMessage = ((McpSchema.TextContent) response.content().get(0)).text(); - assertThat(errorMessage).isEqualTo( - "Response missing structured content which is expected when calling tool with non-empty outputSchema"); - } + HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpServerTransportProvider.getRouterFunction()); + ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler); + this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow(); - mcpServer.close(); + prepareClients(PORT, null); } - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputRuntimeToolAddition(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Start server without tools - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Initially no tools - assertThat(mcpClient.listTools().tools()).isEmpty(); - - // Add tool with output schema at runtime - Map outputSchema = Map.of("type", "object", "properties", - Map.of("message", Map.of("type", "string"), "count", Map.of("type", "integer")), "required", - List.of("message", "count")); - - Tool dynamicTool = Tool.builder() - .name("dynamic-tool") - .description("Dynamically added tool") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification toolSpec = new McpServerFeatures.SyncToolSpecification(dynamicTool, - (exchange, request) -> { - int count = (Integer) request.getOrDefault("count", 1); - return CallToolResult.builder() - .addTextContent("Dynamic tool executed " + count + " times") - .structuredContent(Map.of("message", "Dynamic execution", "count", count)) - .build(); - }); - - // Add tool to server - mcpServer.addTool(toolSpec); - - // Wait for tool list change notification - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(mcpClient.listTools().tools()).hasSize(1); - }); - - // Verify tool was added with output schema - var toolsList = mcpClient.listTools(); - assertThat(toolsList.tools()).hasSize(1); - assertThat(toolsList.tools().get(0).name()).isEqualTo("dynamic-tool"); - // Note: outputSchema might be null in sync server, but validation still works - - // Call dynamically added tool - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("dynamic-tool", Map.of("count", 3))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isFalse(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) response.content().get(0)).text()) - .isEqualTo("Dynamic tool executed 3 times"); - - assertThat(response.structuredContent()).isNotNull(); - assertThatJson(response.structuredContent()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"count":3,"message":"Dynamic execution"}""")); + @AfterEach + public void after() { + if (httpServer != null) { + httpServer.disposeNow(); } - - mcpServer.close(); - } - - private double evaluateExpression(String expression) { - // Simple expression evaluator for testing - return switch (expression) { - case "2 + 3" -> 5.0; - case "10 * 2" -> 20.0; - case "7 + 8" -> 15.0; - case "5 + 3" -> 8.0; - default -> 0.0; - }; } } diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStatelessIntegrationTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStatelessIntegrationTests.java index 2f1765df7..5516e55b7 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStatelessIntegrationTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStatelessIntegrationTests.java @@ -1,53 +1,34 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol; +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; +import org.springframework.web.reactive.function.server.RouterFunctions; + import com.fasterxml.jackson.databind.ObjectMapper; + import io.modelcontextprotocol.client.McpClient; import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport; import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport; import io.modelcontextprotocol.server.McpServer; -import io.modelcontextprotocol.server.McpStatelessServerFeatures; +import io.modelcontextprotocol.server.McpServer.StatelessAsyncSpecification; +import io.modelcontextprotocol.server.McpServer.StatelessSyncSpecification; import io.modelcontextprotocol.server.TestUtil; import io.modelcontextprotocol.server.transport.WebFluxStatelessServerTransport; -import io.modelcontextprotocol.spec.McpSchema; -import io.modelcontextprotocol.spec.McpSchema.CallToolResult; -import io.modelcontextprotocol.spec.McpSchema.CompleteRequest; -import io.modelcontextprotocol.spec.McpSchema.CompleteResult; -import io.modelcontextprotocol.spec.McpSchema.InitializeResult; -import io.modelcontextprotocol.spec.McpSchema.Prompt; -import io.modelcontextprotocol.spec.McpSchema.PromptArgument; -import io.modelcontextprotocol.spec.McpSchema.PromptReference; -import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; -import io.modelcontextprotocol.spec.McpSchema.Tool; -import io.modelcontextprotocol.server.McpTransportContext; -import net.javacrumbs.jsonunit.core.Option; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; -import org.springframework.http.server.reactive.HttpHandler; -import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; -import org.springframework.web.client.RestClient; -import org.springframework.web.reactive.function.client.WebClient; -import org.springframework.web.reactive.function.server.RouterFunctions; import reactor.netty.DisposableServer; import reactor.netty.http.server.HttpServer; -import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.BiFunction; - -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json; -import static org.assertj.core.api.Assertions.assertThat; -import static org.awaitility.Awaitility.await; - -class WebFluxStatelessIntegrationTests { +@Timeout(15) +class WebFluxStatelessIntegrationTests extends AbstractStatelessIntegrationTests { private static final int PORT = TestUtil.findAvailablePort(); @@ -57,19 +38,8 @@ class WebFluxStatelessIntegrationTests { private WebFluxStatelessServerTransport mcpStreamableServerTransport; - ConcurrentHashMap clientBuilders = new ConcurrentHashMap<>(); - - @BeforeEach - public void before() { - this.mcpStreamableServerTransport = WebFluxStatelessServerTransport.builder() - .objectMapper(new ObjectMapper()) - .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) - .build(); - - HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpStreamableServerTransport.getRouterFunction()); - ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler); - this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow(); - + @Override + protected void prepareClients(int port, String mcpEndpoint) { clientBuilders .put("httpclient", McpClient.sync(HttpClientStreamableHttpTransport.builder("http://localhost:" + PORT) @@ -82,391 +52,37 @@ public void before() { .build()) .initializationTimeout(Duration.ofHours(10)) .requestTimeout(Duration.ofHours(10))); - - } - - @AfterEach - public void after() { - if (httpServer != null) { - httpServer.disposeNow(); - } } - // --------------------------------------- - // Tools Tests - // --------------------------------------- - - String emptyJsonSchema = """ - { - "$schema": "http://json-schema.org/draft-07/schema#", - "type": "object", - "properties": {} - } - """; - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testToolCallSuccess(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - McpStatelessServerFeatures.SyncToolSpecification tool1 = new McpStatelessServerFeatures.SyncToolSpecification( - new Tool("tool1", "tool1 description", emptyJsonSchema), (transportContext, request) -> { - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - return callResponse; - }); - - var mcpServer = McpServer.sync(mcpStreamableServerTransport) - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool1) - .build(); - - try (var mcpClient = clientBuilder.build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - } - - mcpServer.close(); + @Override + protected StatelessAsyncSpecification prepareAsyncServerBuilder() { + return McpServer.async(this.mcpStreamableServerTransport); } - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testInitialize(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var mcpServer = McpServer.sync(mcpStreamableServerTransport).build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - } - - mcpServer.close(); + @Override + protected StatelessSyncSpecification prepareSyncServerBuilder() { + return McpServer.sync(this.mcpStreamableServerTransport); } - // --------------------------------------- - // Completion Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : Completion call") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCompletionShouldReturnExpectedSuggestions(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - var expectedValues = List.of("python", "pytorch", "pyside"); - var completionResponse = new CompleteResult(new CompleteResult.CompleteCompletion(expectedValues, 10, // total - true // hasMore - )); - - AtomicReference samplingRequest = new AtomicReference<>(); - BiFunction completionHandler = (transportContext, - request) -> { - samplingRequest.set(request); - return completionResponse; - }; - - var mcpServer = McpServer.sync(mcpStreamableServerTransport) - .capabilities(ServerCapabilities.builder().completions().build()) - .prompts(new McpStatelessServerFeatures.SyncPromptSpecification( - new Prompt("code_review", "Code review", "this is code review prompt", - List.of(new PromptArgument("language", "Language", "string", false))), - (transportContext, getPromptRequest) -> null)) - .completions(new McpStatelessServerFeatures.SyncCompletionSpecification( - new PromptReference("ref/prompt", "code_review", "Code review"), completionHandler)) - .build(); - - try (var mcpClient = clientBuilder.build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CompleteRequest request = new CompleteRequest( - new PromptReference("ref/prompt", "code_review", "Code review"), - new CompleteRequest.CompleteArgument("language", "py")); - - CompleteResult result = mcpClient.completeCompletion(request); - - assertThat(result).isNotNull(); - - assertThat(samplingRequest.get().argument().name()).isEqualTo("language"); - assertThat(samplingRequest.get().argument().value()).isEqualTo("py"); - assertThat(samplingRequest.get().ref().type()).isEqualTo("ref/prompt"); - } - - mcpServer.close(); - } - - // --------------------------------------- - // Tool Structured Output Schema Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputValidationSuccess(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create a tool with output schema - Map outputSchema = Map.of( - "type", "object", "properties", Map.of("result", Map.of("type", "number"), "operation", - Map.of("type", "string"), "timestamp", Map.of("type", "string")), - "required", List.of("result", "operation")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpStatelessServerFeatures.SyncToolSpecification tool = new McpStatelessServerFeatures.SyncToolSpecification( - calculatorTool, (transportContext, request) -> { - String expression = (String) request.arguments().getOrDefault("expression", "2 + 3"); - double result = evaluateExpression(expression); - return CallToolResult.builder() - .structuredContent( - Map.of("result", result, "operation", expression, "timestamp", "2024-01-01T10:00:00Z")) - .build(); - }); - - var mcpServer = McpServer.sync(mcpStreamableServerTransport) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Verify tool is listed with output schema - var toolsList = mcpClient.listTools(); - assertThat(toolsList.tools()).hasSize(1); - assertThat(toolsList.tools().get(0).name()).isEqualTo("calculator"); - // Note: outputSchema might be null in sync server, but validation still works - - // Call tool with valid structured output - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isFalse(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - assertThatJson(((McpSchema.TextContent) response.content().get(0)).text()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"result":5.0,"operation":"2 + 3","timestamp":"2024-01-01T10:00:00Z"}""")); - - assertThat(response.structuredContent()).isNotNull(); - assertThatJson(response.structuredContent()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"result":5.0,"operation":"2 + 3","timestamp":"2024-01-01T10:00:00Z"}""")); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputValidationFailure(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create a tool with output schema - Map outputSchema = Map.of("type", "object", "properties", - Map.of("result", Map.of("type", "number"), "operation", Map.of("type", "string")), "required", - List.of("result", "operation")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpStatelessServerFeatures.SyncToolSpecification tool = new McpStatelessServerFeatures.SyncToolSpecification( - calculatorTool, (transportContext, request) -> { - // Return invalid structured output. Result should be number, missing - // operation - return CallToolResult.builder() - .addTextContent("Invalid calculation") - .structuredContent(Map.of("result", "not-a-number", "extra", "field")) - .build(); - }); - - var mcpServer = McpServer.sync(mcpStreamableServerTransport) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call tool with invalid structured output - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isTrue(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - String errorMessage = ((McpSchema.TextContent) response.content().get(0)).text(); - assertThat(errorMessage).contains("Validation failed"); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputMissingStructuredContent(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create a tool with output schema - Map outputSchema = Map.of("type", "object", "properties", - Map.of("result", Map.of("type", "number")), "required", List.of("result")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpStatelessServerFeatures.SyncToolSpecification tool = new McpStatelessServerFeatures.SyncToolSpecification( - calculatorTool, (transportContext, request) -> { - // Return result without structured content but tool has output schema - return CallToolResult.builder().addTextContent("Calculation completed").build(); - }); - - var mcpServer = McpServer.sync(mcpStreamableServerTransport) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .instructions("bla") - .tools(tool) + @BeforeEach + public void before() { + this.mcpStreamableServerTransport = WebFluxStatelessServerTransport.builder() + .objectMapper(new ObjectMapper()) + .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) .build(); - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call tool that should return structured content but doesn't - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isTrue(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - String errorMessage = ((McpSchema.TextContent) response.content().get(0)).text(); - assertThat(errorMessage).isEqualTo( - "Response missing structured content which is expected when calling tool with non-empty outputSchema"); - } + HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpStreamableServerTransport.getRouterFunction()); + ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler); + this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow(); - mcpServer.close(); + prepareClients(PORT, null); } - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputRuntimeToolAddition(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Start server without tools - var mcpServer = McpServer.sync(mcpStreamableServerTransport) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Initially no tools - assertThat(mcpClient.listTools().tools()).isEmpty(); - - // Add tool with output schema at runtime - Map outputSchema = Map.of("type", "object", "properties", - Map.of("message", Map.of("type", "string"), "count", Map.of("type", "integer")), "required", - List.of("message", "count")); - - Tool dynamicTool = Tool.builder() - .name("dynamic-tool") - .description("Dynamically added tool") - .outputSchema(outputSchema) - .build(); - - McpStatelessServerFeatures.SyncToolSpecification toolSpec = new McpStatelessServerFeatures.SyncToolSpecification( - dynamicTool, (transportContext, request) -> { - int count = (Integer) request.arguments().getOrDefault("count", 1); - return CallToolResult.builder() - .addTextContent("Dynamic tool executed " + count + " times") - .structuredContent(Map.of("message", "Dynamic execution", "count", count)) - .build(); - }); - - // Add tool to server - mcpServer.addTool(toolSpec); - - // Wait for tool list change notification - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(mcpClient.listTools().tools()).hasSize(1); - }); - - // Verify tool was added with output schema - var toolsList = mcpClient.listTools(); - assertThat(toolsList.tools()).hasSize(1); - assertThat(toolsList.tools().get(0).name()).isEqualTo("dynamic-tool"); - // Note: outputSchema might be null in sync server, but validation still works - - // Call dynamically added tool - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("dynamic-tool", Map.of("count", 3))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isFalse(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) response.content().get(0)).text()) - .isEqualTo("Dynamic tool executed 3 times"); - - assertThat(response.structuredContent()).isNotNull(); - assertThatJson(response.structuredContent()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"count":3,"message":"Dynamic execution"}""")); + @AfterEach + public void after() { + if (httpServer != null) { + httpServer.disposeNow(); } - - mcpServer.close(); - } - - private double evaluateExpression(String expression) { - // Simple expression evaluator for testing - return switch (expression) { - case "2 + 3" -> 5.0; - case "10 * 2" -> 20.0; - case "7 + 8" -> 15.0; - case "5 + 3" -> 8.0; - default -> 0.0; - }; } } diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStreamableIntegrationTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStreamableIntegrationTests.java index bc13ad9c6..a7aac0f1e 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStreamableIntegrationTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/WebFluxStreamableIntegrationTests.java @@ -1,72 +1,36 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol; +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; +import org.springframework.web.reactive.function.server.RouterFunctions; +import org.springframework.web.reactive.function.server.ServerRequest; + import com.fasterxml.jackson.databind.ObjectMapper; + import io.modelcontextprotocol.client.McpClient; import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport; import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport; import io.modelcontextprotocol.server.McpServer; -import io.modelcontextprotocol.server.McpServerFeatures; -import io.modelcontextprotocol.server.McpSyncServerExchange; +import io.modelcontextprotocol.server.McpServer.AsyncSpecification; +import io.modelcontextprotocol.server.McpServer.SyncSpecification; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.server.TestUtil; import io.modelcontextprotocol.server.transport.WebFluxStreamableServerTransportProvider; -import io.modelcontextprotocol.spec.McpError; -import io.modelcontextprotocol.spec.McpSchema; -import io.modelcontextprotocol.spec.McpSchema.CallToolResult; -import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities; -import io.modelcontextprotocol.spec.McpSchema.CompleteRequest; -import io.modelcontextprotocol.spec.McpSchema.CompleteResult; -import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest; -import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult; -import io.modelcontextprotocol.spec.McpSchema.ElicitRequest; -import io.modelcontextprotocol.spec.McpSchema.ElicitResult; -import io.modelcontextprotocol.spec.McpSchema.InitializeResult; -import io.modelcontextprotocol.spec.McpSchema.ModelPreferences; -import io.modelcontextprotocol.spec.McpSchema.Prompt; -import io.modelcontextprotocol.spec.McpSchema.PromptArgument; -import io.modelcontextprotocol.spec.McpSchema.PromptReference; -import io.modelcontextprotocol.spec.McpSchema.Role; -import io.modelcontextprotocol.spec.McpSchema.Root; -import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; -import io.modelcontextprotocol.spec.McpSchema.Tool; -import net.javacrumbs.jsonunit.core.Option; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; -import org.springframework.http.server.reactive.HttpHandler; -import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; -import org.springframework.web.client.RestClient; -import org.springframework.web.reactive.function.client.WebClient; -import org.springframework.web.reactive.function.server.RouterFunctions; -import reactor.core.publisher.Mono; import reactor.netty.DisposableServer; import reactor.netty.http.server.HttpServer; -import reactor.test.StepVerifier; -import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.BiFunction; -import java.util.function.Function; -import java.util.stream.Collectors; - -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json; -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatExceptionOfType; -import static org.assertj.core.api.Assertions.assertWith; -import static org.awaitility.Awaitility.await; -import static org.mockito.Mockito.mock; - -class WebFluxStreamableIntegrationTests { +@Timeout(15) +class WebFluxStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests { private static final int PORT = TestUtil.findAvailablePort(); @@ -76,7 +40,37 @@ class WebFluxStreamableIntegrationTests { private WebFluxStreamableServerTransportProvider mcpStreamableServerTransportProvider; - ConcurrentHashMap clientBuilders = new ConcurrentHashMap<>(); + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + + @Override + protected void prepareClients(int port, String mcpEndpoint) { + + clientBuilders + .put("httpclient", + McpClient.sync(HttpClientStreamableHttpTransport.builder("http://localhost:" + PORT) + .endpoint(CUSTOM_MESSAGE_ENDPOINT) + .build()).requestTimeout(Duration.ofHours(10))); + clientBuilders.put("webflux", + McpClient + .sync(WebClientStreamableHttpTransport + .builder(WebClient.builder().baseUrl("http://localhost:" + PORT)) + .endpoint(CUSTOM_MESSAGE_ENDPOINT) + .build()) + .requestTimeout(Duration.ofHours(10))); + } + + @Override + protected AsyncSpecification prepareAsyncServerBuilder() { + return McpServer.async(mcpStreamableServerTransportProvider); + } + + @Override + protected SyncSpecification prepareSyncServerBuilder() { + return McpServer.sync(mcpStreamableServerTransportProvider); + } @BeforeEach public void before() { @@ -84,6 +78,7 @@ public void before() { this.mcpStreamableServerTransportProvider = WebFluxStreamableServerTransportProvider.builder() .objectMapper(new ObjectMapper()) .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .build(); HttpHandler httpHandler = RouterFunctions @@ -91,19 +86,7 @@ public void before() { ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler); this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow(); - clientBuilders - .put("httpclient", - McpClient.sync(HttpClientStreamableHttpTransport.builder("http://localhost:" + PORT) - .endpoint(CUSTOM_MESSAGE_ENDPOINT) - .build()).initializationTimeout(Duration.ofHours(10)).requestTimeout(Duration.ofHours(10))); - clientBuilders - .put("webflux", McpClient - .sync(WebClientStreamableHttpTransport.builder(WebClient.builder().baseUrl("http://localhost:" + PORT)) - .endpoint(CUSTOM_MESSAGE_ENDPOINT) - .build()) - .initializationTimeout(Duration.ofHours(10)) - .requestTimeout(Duration.ofHours(10))); - + prepareClients(PORT, null); } @AfterEach @@ -113,1380 +96,4 @@ public void after() { } } - // --------------------------------------- - // Sampling Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateMessageWithoutSamplingCapabilities(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> exchange.createMessage(mock(CreateMessageRequest.class)) - .thenReturn(mock(CallToolResult.class))) - .build(); - - var server = McpServer.async(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample " + "client", "0.0.0")) - .build();) { - - assertThat(client.initialize()).isNotNull(); - - try { - client.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - } - catch (McpError e) { - assertThat(e).isInstanceOf(McpError.class) - .hasMessage("Client must be configured with sampling capabilities"); - } - } - server.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateMessageSuccess(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - Function samplingHandler = request -> { - assertThat(request.messages()).hasSize(1); - assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); - - return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", - CreateMessageResult.StopReason.STOP_SEQUENCE); - }; - - CallToolResult callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - - AtomicReference samplingResult = new AtomicReference<>(); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var createMessageRequest = CreateMessageRequest.builder() - .messages(List - .of(new McpSchema.SamplingMessage(Role.USER, new McpSchema.TextContent("Test message")))) - .modelPreferences(ModelPreferences.builder() - .hints(List.of()) - .costPriority(1.0) - .speedPriority(1.0) - .intelligencePriority(1.0) - .build()) - .build(); - - return exchange.createMessage(createMessageRequest) - .doOnNext(samplingResult::set) - .thenReturn(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - - assertWith(samplingResult.get(), result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }); - } - mcpServer.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws InterruptedException { - - // Client - var clientBuilder = clientBuilders.get(clientType); - - Function samplingHandler = request -> { - assertThat(request.messages()).hasSize(1); - assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", - CreateMessageResult.StopReason.STOP_SEQUENCE); - }; - - // Server - - CallToolResult callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - - AtomicReference samplingResult = new AtomicReference<>(); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var craeteMessageRequest = CreateMessageRequest.builder() - .messages(List - .of(new McpSchema.SamplingMessage(Role.USER, new McpSchema.TextContent("Test message")))) - .modelPreferences(ModelPreferences.builder() - .hints(List.of()) - .costPriority(1.0) - .speedPriority(1.0) - .intelligencePriority(1.0) - .build()) - .build(); - - return exchange.createMessage(craeteMessageRequest) - .doOnNext(samplingResult::set) - .thenReturn(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpStreamableServerTransportProvider) - .requestTimeout(Duration.ofSeconds(4)) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - - assertWith(samplingResult.get(), result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }); - } - - mcpServer.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateMessageWithRequestTimeoutFail(String clientType) throws InterruptedException { - - // Client - var clientBuilder = clientBuilders.get(clientType); - - Function samplingHandler = request -> { - assertThat(request.messages()).hasSize(1); - assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", - CreateMessageResult.StopReason.STOP_SEQUENCE); - }; - - // Server - - CallToolResult callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var craeteMessageRequest = CreateMessageRequest.builder() - .messages(List - .of(new McpSchema.SamplingMessage(Role.USER, new McpSchema.TextContent("Test message")))) - .build(); - - return exchange.createMessage(craeteMessageRequest).thenReturn(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpStreamableServerTransportProvider) - .requestTimeout(Duration.ofSeconds(1)) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThatExceptionOfType(McpError.class).isThrownBy(() -> { - mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - }).withMessageContaining("within 1000ms"); - - } - - mcpServer.closeGracefully().block(); - } - - // --------------------------------------- - // Elicitation Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateElicitationWithoutElicitationCapabilities(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> exchange.createElicitation(mock(ElicitRequest.class)) - .then(Mono.just(mock(CallToolResult.class)))) - .build(); - - var server = McpServer.async(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try ( - // Create client without elicitation capabilities - var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) { - - assertThat(client.initialize()).isNotNull(); - - try { - client.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - } - catch (McpError e) { - assertThat(e).isInstanceOf(McpError.class) - .hasMessage("Client must be configured with elicitation capabilities"); - } - } - server.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateElicitationSuccess(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - Function elicitationHandler = request -> { - assertThat(request.message()).isNotEmpty(); - assertThat(request.requestedSchema()).isNotNull(); - - return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); - }; - - CallToolResult callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var elicitationRequest = ElicitRequest.builder() - .message("Test message") - .requestedSchema( - Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) - .build(); - - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().elicitation().build()) - .elicitation(elicitationHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - } - mcpServer.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { - - // Client - var clientBuilder = clientBuilders.get(clientType); - - Function elicitationHandler = request -> { - assertThat(request.message()).isNotEmpty(); - assertThat(request.requestedSchema()).isNotNull(); - return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); - }; - - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().elicitation().build()) - .elicitation(elicitationHandler) - .build(); - - // Server - - CallToolResult callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var elicitationRequest = ElicitRequest.builder() - .message("Test message") - .requestedSchema( - Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) - .build(); - - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(3)) - .tools(tool) - .build(); - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - - mcpClient.closeGracefully(); - mcpServer.closeGracefully().block(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCreateElicitationWithRequestTimeoutFail(String clientType) { - - var latch = new CountDownLatch(1); - // Client - var clientBuilder = clientBuilders.get(clientType); - - Function elicitationHandler = request -> { - assertThat(request.message()).isNotEmpty(); - assertThat(request.requestedSchema()).isNotNull(); - - try { - if (!latch.await(2, TimeUnit.SECONDS)) { - throw new RuntimeException("Timeout waiting for elicitation processing"); - } - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); - }; - - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().elicitation().build()) - .elicitation(elicitationHandler) - .build(); - - // Server - - CallToolResult callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - - AtomicReference resultRef = new AtomicReference<>(); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var elicitationRequest = ElicitRequest.builder() - .message("Test message") - .requestedSchema( - Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) - .build(); - - return exchange.createElicitation(elicitationRequest) - .doOnNext(resultRef::set) - .then(Mono.just(callResponse)); - }) - .build(); - - var mcpServer = McpServer.async(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(1)) // 1 second. - .tools(tool) - .build(); - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThatExceptionOfType(McpError.class).isThrownBy(() -> { - mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - }).withMessageContaining("within 1000ms"); - - ElicitResult elicitResult = resultRef.get(); - assertThat(elicitResult).isNull(); - - mcpClient.closeGracefully(); - mcpServer.closeGracefully().block(); - } - - // --------------------------------------- - // Roots Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsSuccess(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - List roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2")); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(roots) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(rootsRef.get()).isNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(roots); - }); - - // Remove a root - mcpClient.removeRoot(roots.get(0).uri()); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(roots.get(1))); - }); - - // Add a new root - var root3 = new Root("uri3://", "root3"); - mcpClient.addRoot(root3); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(roots.get(1), root3)); - }); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsWithoutCapability(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - McpServerFeatures.SyncToolSpecification tool = McpServerFeatures.SyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - exchange.listRoots(); // try to list roots - - return mock(CallToolResult.class); - }) - .build(); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> { - }) - .tools(tool) - .build(); - - // Create client without roots capability - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()).build()) { - - assertThat(mcpClient.initialize()).isNotNull(); - - // Attempt to list roots should fail - try { - mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - } - catch (McpError e) { - assertThat(e).isInstanceOf(McpError.class).hasMessage("Roots not supported"); - } - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsNotificationWithEmptyRootsList(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(List.of()) // Empty roots list - .build()) { - - assertThat(mcpClient.initialize()).isNotNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).isEmpty(); - }); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsWithMultipleHandlers(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - List roots = List.of(new Root("uri1://", "root1")); - - AtomicReference> rootsRef1 = new AtomicReference<>(); - AtomicReference> rootsRef2 = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef1.set(rootsUpdate)) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef2.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(roots) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef1.get()).containsAll(roots); - assertThat(rootsRef2.get()).containsAll(roots); - }); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testRootsServerCloseWithActiveSubscription(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - List roots = List.of(new Root("uri1://", "root1")); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(roots) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(roots); - }); - } - - mcpServer.close(); - } - - // --------------------------------------- - // Tools Tests - // --------------------------------------- - - String emptyJsonSchema = """ - { - "$schema": "http://json-schema.org/draft-07/schema#", - "type": "object", - "properties": {} - } - """; - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testToolCallSuccess(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - return callResponse; - }) - .build(); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool1) - .build(); - - try (var mcpClient = clientBuilder.build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testToolListChangeHandlingSuccess(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var callResponse = new CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() - .tool(new Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - return callResponse; - }) - .build(); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool1) - .build(); - - try (var mcpClient = clientBuilder.toolsChangeConsumer(toolsUpdate -> { - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - rootsRef.set(toolsUpdate); - }).build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(rootsRef.get()).isNull(); - - assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); - - mcpServer.notifyToolsListChanged(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool1.tool())); - }); - - // Remove a tool - mcpServer.removeTool("tool1"); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).isEmpty(); - }); - - // Add a new tool - McpServerFeatures.SyncToolSpecification tool2 = McpServerFeatures.SyncToolSpecification.builder() - .tool(new Tool("tool2", "tool2 description", emptyJsonSchema)) - .callHandler((exchange, request) -> callResponse) - .build(); - - mcpServer.addTool(tool2); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool2.tool())); - }); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testInitialize(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider).build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - } - - mcpServer.close(); - } - - // --------------------------------------- - // Logging Tests - // --------------------------------------- - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testLoggingNotification(String clientType) throws InterruptedException { - int expectedNotificationsCount = 3; - CountDownLatch latch = new CountDownLatch(expectedNotificationsCount); - // Create a list to store received logging notifications - List receivedNotifications = new CopyOnWriteArrayList<>(); - - var clientBuilder = clientBuilders.get(clientType); - - // Create server with a tool that sends logging notifications - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("logging-test", "Test logging notifications", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - // Create and send notifications with different levels - - //@formatter:off - return exchange // This should be filtered out (DEBUG < NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.DEBUG) - .logger("test-logger") - .data("Debug message") - .build()) - .then(exchange // This should be sent (NOTICE >= NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.NOTICE) - .logger("test-logger") - .data("Notice message") - .build())) - .then(exchange // This should be sent (ERROR > NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.ERROR) - .logger("test-logger") - .data("Error message") - .build())) - .then(exchange // This should be filtered out (INFO < NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.INFO) - .logger("test-logger") - .data("Another info message") - .build())) - .then(exchange // This should be sent (ERROR >= NOTICE) - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.ERROR) - .logger("test-logger") - .data("Another error message") - .build())) - .thenReturn(new CallToolResult("Logging test completed", false)); - //@formatter:on - }) - .build(); - - var mcpServer = McpServer.async(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().logging().tools(true).build()) - .tools(tool) - .build(); - - try ( - // Create client with logging notification handler - var mcpClient = clientBuilder.loggingConsumer(notification -> { - receivedNotifications.add(notification); - latch.countDown(); - }).build()) { - - // Initialize client - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Set minimum logging level to NOTICE - mcpClient.setLoggingLevel(McpSchema.LoggingLevel.NOTICE); - - // Call the tool that sends logging notifications - CallToolResult result = mcpClient.callTool(new McpSchema.CallToolRequest("logging-test", Map.of())); - assertThat(result).isNotNull(); - assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Logging test completed"); - - assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue(); - - // Should have received 3 notifications (1 NOTICE and 2 ERROR) - assertThat(receivedNotifications).hasSize(expectedNotificationsCount); - - Map notificationMap = receivedNotifications.stream() - .collect(Collectors.toMap(n -> n.data(), n -> n)); - - // First notification should be NOTICE level - assertThat(notificationMap.get("Notice message").level()).isEqualTo(McpSchema.LoggingLevel.NOTICE); - assertThat(notificationMap.get("Notice message").logger()).isEqualTo("test-logger"); - assertThat(notificationMap.get("Notice message").data()).isEqualTo("Notice message"); - - // Second notification should be ERROR level - assertThat(notificationMap.get("Error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); - assertThat(notificationMap.get("Error message").logger()).isEqualTo("test-logger"); - assertThat(notificationMap.get("Error message").data()).isEqualTo("Error message"); - - // Third notification should be ERROR level - assertThat(notificationMap.get("Another error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); - assertThat(notificationMap.get("Another error message").logger()).isEqualTo("test-logger"); - assertThat(notificationMap.get("Another error message").data()).isEqualTo("Another error message"); - } - mcpServer.close(); - } - - // --------------------------------------- - // Progress Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testProgressNotification(String clientType) throws InterruptedException { - int expectedNotificationsCount = 4; // 3 notifications + 1 for another progress - // token - CountDownLatch latch = new CountDownLatch(expectedNotificationsCount); - // Create a list to store received logging notifications - List receivedNotifications = new CopyOnWriteArrayList<>(); - - var clientBuilder = clientBuilders.get(clientType); - - // Create server with a tool that sends logging notifications - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(Tool.builder() - .name("progress-test") - .description("Test progress notifications") - .inputSchema(emptyJsonSchema) - .build()) - .callHandler((exchange, request) -> { - - // Create and send notifications - var progressToken = (String) request.meta().get("progressToken"); - - return exchange - .progressNotification( - new McpSchema.ProgressNotification(progressToken, 0.0, 1.0, "Processing started")) - .then(exchange.progressNotification( - new McpSchema.ProgressNotification(progressToken, 0.5, 1.0, "Processing data"))) - .then(// Send a progress notification with another progress value - // should - exchange.progressNotification(new McpSchema.ProgressNotification("another-progress-token", - 0.0, 1.0, "Another processing started"))) - .then(exchange.progressNotification( - new McpSchema.ProgressNotification(progressToken, 1.0, 1.0, "Processing completed"))) - .thenReturn(new CallToolResult(("Progress test completed"), false)); - }) - .build(); - - var mcpServer = McpServer.async(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try ( - // Create client with progress notification handler - var mcpClient = clientBuilder.progressConsumer(notification -> { - receivedNotifications.add(notification); - latch.countDown(); - }).build()) { - - // Initialize client - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call the tool that sends progress notifications - McpSchema.CallToolRequest callToolRequest = McpSchema.CallToolRequest.builder() - .name("progress-test") - .meta(Map.of("progressToken", "test-progress-token")) - .build(); - CallToolResult result = mcpClient.callTool(callToolRequest); - assertThat(result).isNotNull(); - assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Progress test completed"); - - assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue(); - - // Should have received 3 notifications - assertThat(receivedNotifications).hasSize(expectedNotificationsCount); - - Map notificationMap = receivedNotifications.stream() - .collect(Collectors.toMap(n -> n.message(), n -> n)); - - // First notification should be 0.0/1.0 progress - assertThat(notificationMap.get("Processing started").progressToken()).isEqualTo("test-progress-token"); - assertThat(notificationMap.get("Processing started").progress()).isEqualTo(0.0); - assertThat(notificationMap.get("Processing started").total()).isEqualTo(1.0); - assertThat(notificationMap.get("Processing started").message()).isEqualTo("Processing started"); - - // Second notification should be 0.5/1.0 progress - assertThat(notificationMap.get("Processing data").progressToken()).isEqualTo("test-progress-token"); - assertThat(notificationMap.get("Processing data").progress()).isEqualTo(0.5); - assertThat(notificationMap.get("Processing data").total()).isEqualTo(1.0); - assertThat(notificationMap.get("Processing data").message()).isEqualTo("Processing data"); - - // Third notification should be another progress token with 0.0/1.0 progress - assertThat(notificationMap.get("Another processing started").progressToken()) - .isEqualTo("another-progress-token"); - assertThat(notificationMap.get("Another processing started").progress()).isEqualTo(0.0); - assertThat(notificationMap.get("Another processing started").total()).isEqualTo(1.0); - assertThat(notificationMap.get("Another processing started").message()) - .isEqualTo("Another processing started"); - - // Fourth notification should be 1.0/1.0 progress - assertThat(notificationMap.get("Processing completed").progressToken()).isEqualTo("test-progress-token"); - assertThat(notificationMap.get("Processing completed").progress()).isEqualTo(1.0); - assertThat(notificationMap.get("Processing completed").total()).isEqualTo(1.0); - assertThat(notificationMap.get("Processing completed").message()).isEqualTo("Processing completed"); - } - finally { - mcpServer.close(); - } - } - - // --------------------------------------- - // Completion Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : Completion call") - @ValueSource(strings = { "httpclient", "webflux" }) - void testCompletionShouldReturnExpectedSuggestions(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - var expectedValues = List.of("python", "pytorch", "pyside"); - var completionResponse = new CompleteResult(new CompleteResult.CompleteCompletion(expectedValues, 10, // total - true // hasMore - )); - - AtomicReference samplingRequest = new AtomicReference<>(); - BiFunction completionHandler = (mcpSyncServerExchange, - request) -> { - samplingRequest.set(request); - return completionResponse; - }; - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .capabilities(ServerCapabilities.builder().completions().build()) - .prompts(new McpServerFeatures.SyncPromptSpecification( - new Prompt("code_review", "Code review", "this is code review prompt", - List.of(new PromptArgument("language", "Language", "string", false))), - (mcpSyncServerExchange, getPromptRequest) -> null)) - .completions(new McpServerFeatures.SyncCompletionSpecification( - new PromptReference("ref/prompt", "code_review", "Code review"), completionHandler)) - .build(); - - try (var mcpClient = clientBuilder.build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CompleteRequest request = new CompleteRequest( - new PromptReference("ref/prompt", "code_review", "Code review"), - new CompleteRequest.CompleteArgument("language", "py")); - - CompleteResult result = mcpClient.completeCompletion(request); - - assertThat(result).isNotNull(); - - assertThat(samplingRequest.get().argument().name()).isEqualTo("language"); - assertThat(samplingRequest.get().argument().value()).isEqualTo("py"); - assertThat(samplingRequest.get().ref().type()).isEqualTo("ref/prompt"); - } - - mcpServer.close(); - } - - // --------------------------------------- - // Ping Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testPingSuccess(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create server with a tool that uses ping functionality - AtomicReference executionOrder = new AtomicReference<>(""); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new Tool("ping-async-test", "Test ping async behavior", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - executionOrder.set(executionOrder.get() + "1"); - - // Test async ping behavior - return exchange.ping().doOnNext(result -> { - - assertThat(result).isNotNull(); - // Ping should return an empty object or map - assertThat(result).isInstanceOf(Map.class); - - executionOrder.set(executionOrder.get() + "2"); - assertThat(result).isNotNull(); - }).then(Mono.fromCallable(() -> { - executionOrder.set(executionOrder.get() + "3"); - return new CallToolResult("Async ping test completed", false); - })); - }) - .build(); - - var mcpServer = McpServer.async(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - - // Initialize client - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call the tool that tests ping async behavior - CallToolResult result = mcpClient.callTool(new McpSchema.CallToolRequest("ping-async-test", Map.of())); - assertThat(result).isNotNull(); - assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Async ping test completed"); - - // Verify execution order - assertThat(executionOrder.get()).isEqualTo("123"); - } - - mcpServer.closeGracefully().block(); - } - - // --------------------------------------- - // Tool Structured Output Schema Tests - // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputValidationSuccess(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create a tool with output schema - Map outputSchema = Map.of( - "type", "object", "properties", Map.of("result", Map.of("type", "number"), "operation", - Map.of("type", "string"), "timestamp", Map.of("type", "string")), - "required", List.of("result", "operation")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification tool = new McpServerFeatures.SyncToolSpecification(calculatorTool, - (exchange, request) -> { - String expression = (String) request.getOrDefault("expression", "2 + 3"); - double result = evaluateExpression(expression); - return CallToolResult.builder() - .structuredContent( - Map.of("result", result, "operation", expression, "timestamp", "2024-01-01T10:00:00Z")) - .build(); - }); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Verify tool is listed with output schema - var toolsList = mcpClient.listTools(); - assertThat(toolsList.tools()).hasSize(1); - assertThat(toolsList.tools().get(0).name()).isEqualTo("calculator"); - // Note: outputSchema might be null in sync server, but validation still works - - // Call tool with valid structured output - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isFalse(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - assertThatJson(((McpSchema.TextContent) response.content().get(0)).text()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"result":5.0,"operation":"2 + 3","timestamp":"2024-01-01T10:00:00Z"}""")); - - assertThat(response.structuredContent()).isNotNull(); - assertThatJson(response.structuredContent()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"result":5.0,"operation":"2 + 3","timestamp":"2024-01-01T10:00:00Z"}""")); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputValidationFailure(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create a tool with output schema - Map outputSchema = Map.of("type", "object", "properties", - Map.of("result", Map.of("type", "number"), "operation", Map.of("type", "string")), "required", - List.of("result", "operation")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification tool = new McpServerFeatures.SyncToolSpecification(calculatorTool, - (exchange, request) -> { - // Return invalid structured output. Result should be number, missing - // operation - return CallToolResult.builder() - .addTextContent("Invalid calculation") - .structuredContent(Map.of("result", "not-a-number", "extra", "field")) - .build(); - }); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call tool with invalid structured output - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isTrue(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - String errorMessage = ((McpSchema.TextContent) response.content().get(0)).text(); - assertThat(errorMessage).contains("Validation failed"); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputMissingStructuredContent(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Create a tool with output schema - Map outputSchema = Map.of("type", "object", "properties", - Map.of("result", Map.of("type", "number")), "required", List.of("result")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification tool = new McpServerFeatures.SyncToolSpecification(calculatorTool, - (exchange, request) -> { - // Return result without structured content but tool has output schema - return CallToolResult.builder().addTextContent("Calculation completed").build(); - }); - - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .instructions("bla") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call tool that should return structured content but doesn't - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isTrue(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - String errorMessage = ((McpSchema.TextContent) response.content().get(0)).text(); - assertThat(errorMessage).isEqualTo( - "Response missing structured content which is expected when calling tool with non-empty outputSchema"); - } - - mcpServer.close(); - } - - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void testStructuredOutputRuntimeToolAddition(String clientType) { - var clientBuilder = clientBuilders.get(clientType); - - // Start server without tools - var mcpServer = McpServer.sync(mcpStreamableServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Initially no tools - assertThat(mcpClient.listTools().tools()).isEmpty(); - - // Add tool with output schema at runtime - Map outputSchema = Map.of("type", "object", "properties", - Map.of("message", Map.of("type", "string"), "count", Map.of("type", "integer")), "required", - List.of("message", "count")); - - Tool dynamicTool = Tool.builder() - .name("dynamic-tool") - .description("Dynamically added tool") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification toolSpec = new McpServerFeatures.SyncToolSpecification(dynamicTool, - (exchange, request) -> { - int count = (Integer) request.getOrDefault("count", 1); - return CallToolResult.builder() - .addTextContent("Dynamic tool executed " + count + " times") - .structuredContent(Map.of("message", "Dynamic execution", "count", count)) - .build(); - }); - - // Add tool to server - mcpServer.addTool(toolSpec); - - // Wait for tool list change notification - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(mcpClient.listTools().tools()).hasSize(1); - }); - - // Verify tool was added with output schema - var toolsList = mcpClient.listTools(); - assertThat(toolsList.tools()).hasSize(1); - assertThat(toolsList.tools().get(0).name()).isEqualTo("dynamic-tool"); - // Note: outputSchema might be null in sync server, but validation still works - - // Call dynamically added tool - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("dynamic-tool", Map.of("count", 3))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isFalse(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) response.content().get(0)).text()) - .isEqualTo("Dynamic tool executed 3 times"); - - assertThat(response.structuredContent()).isNotNull(); - assertThatJson(response.structuredContent()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"count":3,"message":"Dynamic execution"}""")); - } - - mcpServer.close(); - } - - private double evaluateExpression(String expression) { - // Simple expression evaluator for testing - return switch (expression) { - case "2 + 3" -> 5.0; - case "10 * 2" -> 20.0; - case "7 + 8" -> 15.0; - case "5 + 3" -> 8.0; - default -> 0.0; - }; - } - } diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpAsyncClientResiliencyTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpAsyncClientResiliencyTests.java index 7c4d35db8..191f10376 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpAsyncClientResiliencyTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpAsyncClientResiliencyTests.java @@ -1,3 +1,7 @@ +/* + * Copyright 2025-2025 the original author or authors. + */ + package io.modelcontextprotocol.client; import io.modelcontextprotocol.client.transport.WebClientStreamableHttpTransport; diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpAsyncClientTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpAsyncClientTests.java index 5ff707b3c..f8a16c153 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpAsyncClientTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpAsyncClientTests.java @@ -1,3 +1,7 @@ +/* + * Copyright 2025-2025 the original author or authors. + */ + package io.modelcontextprotocol.client; import org.junit.jupiter.api.Timeout; diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpSyncClientTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpSyncClientTests.java index 70260c8bf..5e9960d0e 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpSyncClientTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/WebClientStreamableHttpSyncClientTests.java @@ -1,3 +1,7 @@ +/* + * Copyright 2025-2025 the original author or authors. + */ + package io.modelcontextprotocol.client; import org.junit.jupiter.api.Timeout; diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransportErrorHandlingTest.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransportErrorHandlingTest.java new file mode 100644 index 000000000..cdbb97e17 --- /dev/null +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/client/transport/WebClientStreamableHttpTransportErrorHandlingTest.java @@ -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 serverResponseStatus = new AtomicReference<>(200); + + private AtomicReference currentServerSessionId = new AtomicReference<>(null); + + private AtomicReference 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 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 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 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); + } + +} diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/WebFluxStreamableMcpAsyncServerTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/WebFluxStreamableMcpAsyncServerTests.java index 928bd812d..959f2f472 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/WebFluxStreamableMcpAsyncServerTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/WebFluxStreamableMcpAsyncServerTests.java @@ -5,9 +5,7 @@ package io.modelcontextprotocol.server; import com.fasterxml.jackson.databind.ObjectMapper; -import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider; import io.modelcontextprotocol.server.transport.WebFluxStreamableServerTransportProvider; -import io.modelcontextprotocol.spec.McpServerTransportProvider; import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider; import org.junit.jupiter.api.Timeout; import org.springframework.http.server.reactive.HttpHandler; diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/WebFluxStreamableMcpSyncServerTests.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/WebFluxStreamableMcpSyncServerTests.java index e82e384c4..3396d489c 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/WebFluxStreamableMcpSyncServerTests.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/WebFluxStreamableMcpSyncServerTests.java @@ -5,7 +5,6 @@ package io.modelcontextprotocol.server; import com.fasterxml.jackson.databind.ObjectMapper; -import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider; import io.modelcontextprotocol.server.transport.WebFluxStreamableServerTransportProvider; import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider; import org.junit.jupiter.api.Timeout; diff --git a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/transport/BlockingInputStream.java b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/transport/BlockingInputStream.java index 0ab72a99f..dfb004e9b 100644 --- a/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/transport/BlockingInputStream.java +++ b/mcp-spring/mcp-spring-webflux/src/test/java/io/modelcontextprotocol/server/transport/BlockingInputStream.java @@ -1,6 +1,7 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol.server.transport; import java.io.IOException; diff --git a/mcp-spring/mcp-spring-webmvc/pom.xml b/mcp-spring/mcp-spring-webmvc/pom.xml index 0a18bf711..af9ceb9b7 100644 --- a/mcp-spring/mcp-spring-webmvc/pom.xml +++ b/mcp-spring/mcp-spring-webmvc/pom.xml @@ -6,7 +6,7 @@ io.modelcontextprotocol.sdk mcp-parent - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT ../../pom.xml mcp-spring-webmvc @@ -25,7 +25,7 @@ io.modelcontextprotocol.sdk mcp - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT @@ -37,14 +37,14 @@ io.modelcontextprotocol.sdk mcp-test - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT test io.modelcontextprotocol.sdk mcp-spring-webflux - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT test diff --git a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java index a3898006d..6e92cf10c 100644 --- a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java +++ b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcSseServerTransportProvider.java @@ -6,16 +6,22 @@ import java.io.IOException; import java.time.Duration; +import java.util.List; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.locks.ReentrantLock; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; + +import io.modelcontextprotocol.server.DefaultMcpTransportContext; +import io.modelcontextprotocol.server.McpTransportContext; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpServerTransport; import io.modelcontextprotocol.spec.McpServerTransportProvider; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.spec.McpServerSession; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.KeepAliveScheduler; @@ -104,6 +110,8 @@ public class WebMvcSseServerTransportProvider implements McpServerTransportProvi */ private final ConcurrentHashMap sessions = new ConcurrentHashMap<>(); + private McpTransportContextExtractor contextExtractor; + /** * Flag indicating if the transport is shutting down. */ @@ -175,7 +183,7 @@ public WebMvcSseServerTransportProvider(ObjectMapper objectMapper, String baseUr * messages via HTTP POST. This endpoint will be communicated to clients through the * SSE connection's initial endpoint event. * @param sseEndpoint The endpoint URI where clients establish their SSE connections. - * * @param keepAliveInterval The interval for sending keep-alive messages to + * @param keepAliveInterval The interval for sending keep-alive messages to clients. * @throws IllegalArgumentException if any parameter is null * @deprecated Use the builder {@link #builder()} instead for better configuration * options. @@ -183,15 +191,39 @@ public WebMvcSseServerTransportProvider(ObjectMapper objectMapper, String baseUr @Deprecated public WebMvcSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, String sseEndpoint, Duration keepAliveInterval) { + this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, keepAliveInterval, + (serverRequest, context) -> context); + } + + /** + * Constructs a new WebMvcSseServerTransportProvider instance. + * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization + * of messages. + * @param baseUrl The base URL for the message endpoint, used to construct the full + * endpoint URL for clients. + * @param messageEndpoint The endpoint URI where clients should send their JSON-RPC + * messages via HTTP POST. This endpoint will be communicated to clients through the + * SSE connection's initial endpoint event. + * @param sseEndpoint The endpoint URI where clients establish their SSE connections. + * @param keepAliveInterval The interval for sending keep-alive messages to clients. + * @param contextExtractor The contextExtractor to fill in a + * {@link McpTransportContext}. + * @throws IllegalArgumentException if any parameter is null + */ + private WebMvcSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, + String sseEndpoint, Duration keepAliveInterval, + McpTransportContextExtractor contextExtractor) { Assert.notNull(objectMapper, "ObjectMapper must not be null"); Assert.notNull(baseUrl, "Message base URL must not be null"); Assert.notNull(messageEndpoint, "Message endpoint must not be null"); Assert.notNull(sseEndpoint, "SSE endpoint must not be null"); + Assert.notNull(contextExtractor, "Context extractor must not be null"); this.objectMapper = objectMapper; this.baseUrl = baseUrl; this.messageEndpoint = messageEndpoint; this.sseEndpoint = sseEndpoint; + this.contextExtractor = contextExtractor; this.routerFunction = RouterFunctions.route() .GET(this.sseEndpoint, this::handleSseConnection) .POST(this.messageEndpoint, this::handleMessage) @@ -210,8 +242,8 @@ public WebMvcSseServerTransportProvider(ObjectMapper objectMapper, String baseUr } @Override - public String protocolVersion() { - return "2024-11-05"; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05); } @Override @@ -365,11 +397,17 @@ private ServerResponse handleMessage(ServerRequest request) { } try { + final McpTransportContext transportContext = this.contextExtractor.extract(request, + new DefaultMcpTransportContext()); + String body = request.body(String.class); McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body); // Process the message through the session's handle method - session.handle(message).block(); // Block for WebMVC compatibility + session.handle(message).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)).block(); // Block + // for + // WebMVC + // compatibility return ServerResponse.ok().build(); } @@ -515,6 +553,8 @@ public static class Builder { private Duration keepAliveInterval; + private McpTransportContextExtractor contextExtractor = (serverRequest, context) -> context; + /** * Sets the JSON object mapper to use for message serialization/deserialization. * @param objectMapper The object mapper to use @@ -574,6 +614,22 @@ public Builder keepAliveInterval(Duration keepAliveInterval) { return this; } + /** + * Sets the context extractor that allows providing the MCP feature + * implementations to inspect HTTP transport level metadata that was present at + * HTTP request processing time. This allows to extract custom headers and other + * useful data for use during execution later on in the process. + * @param contextExtractor The contextExtractor to fill in a + * {@link McpTransportContext}. + * @return this builder instance + * @throws IllegalArgumentException if contextExtractor is null + */ + public Builder contextExtractor(McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "contextExtractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + /** * Builds a new instance of WebMvcSseServerTransportProvider with the configured * settings. @@ -585,7 +641,7 @@ public WebMvcSseServerTransportProvider build() { throw new IllegalStateException("MessageEndpoint must be set"); } return new WebMvcSseServerTransportProvider(objectMapper, baseUrl, messageEndpoint, sseEndpoint, - keepAliveInterval); + keepAliveInterval, contextExtractor); } } diff --git a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStatelessServerTransport.java b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStatelessServerTransport.java index 1b026fc46..fef1920fc 100644 --- a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStatelessServerTransport.java +++ b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStatelessServerTransport.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2024 the original author or authors. + */ + package io.modelcontextprotocol.server.transport; import com.fasterxml.jackson.databind.ObjectMapper; diff --git a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java index 2f94d5c11..fa51a0130 100644 --- a/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java +++ b/mcp-spring/mcp-spring-webmvc/src/main/java/io/modelcontextprotocol/server/transport/WebMvcStreamableServerTransportProvider.java @@ -32,6 +32,7 @@ import io.modelcontextprotocol.spec.McpStreamableServerSession; import io.modelcontextprotocol.spec.McpStreamableServerTransport; import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.KeepAliveScheduler; import reactor.core.publisher.Flux; @@ -95,9 +96,6 @@ public class WebMvcStreamableServerTransportProvider implements McpStreamableSer private McpTransportContextExtractor contextExtractor; - // private Function contextExtractor = req -> new - // DefaultMcpTransportContext(); - /** * Flag indicating if the transport is shutting down. */ @@ -142,14 +140,11 @@ private WebMvcStreamableServerTransportProvider(ObjectMapper objectMapper, Strin this.keepAliveScheduler.start(); } - else { - logger.warn("Keep-alive interval is not set or invalid. No keep-alive will be scheduled."); - } } @Override - public String protocolVersion() { - return "2025-03-26"; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05, ProtocolVersions.MCP_2025_03_26); } @Override diff --git a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java index 45f6b94f0..8cb2973ed 100644 --- a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java +++ b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseIntegrationTests.java @@ -11,11 +11,13 @@ 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; import org.springframework.web.servlet.config.annotation.EnableWebMvc; import org.springframework.web.servlet.function.RouterFunction; +import org.springframework.web.servlet.function.ServerRequest; import org.springframework.web.servlet.function.ServerResponse; import com.fasterxml.jackson.databind.ObjectMapper; @@ -29,6 +31,7 @@ import io.modelcontextprotocol.server.transport.WebMvcSseServerTransportProvider; import reactor.core.scheduler.Schedulers; +@Timeout(15) class WebMvcSseIntegrationTests extends AbstractMcpClientServerIntegrationTests { private static final int PORT = TestUtil.findAvailablePort(); @@ -37,16 +40,21 @@ class WebMvcSseIntegrationTests extends AbstractMcpClientServerIntegrationTests private WebMvcSseServerTransportProvider mcpServerTransportProvider; + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + @Override protected void prepareClients(int port, String mcpEndpoint) { clientBuilders.put("httpclient", McpClient.sync(HttpClientSseClientTransport.builder("http://localhost:" + port).build()) - .initializationTimeout(Duration.ofHours(10)) .requestTimeout(Duration.ofHours(10))); clientBuilders.put("webflux", McpClient - .sync(WebFluxSseClientTransport.builder(WebClient.builder().baseUrl("http://localhost:" + port)).build())); + .sync(WebFluxSseClientTransport.builder(WebClient.builder().baseUrl("http://localhost:" + port)).build()) + .requestTimeout(Duration.ofHours(10))); } @Configuration @@ -55,7 +63,11 @@ static class TestConfig { @Bean public WebMvcSseServerTransportProvider webMvcSseServerTransportProvider() { - return new WebMvcSseServerTransportProvider(new ObjectMapper(), MESSAGE_ENDPOINT); + return WebMvcSseServerTransportProvider.builder() + .objectMapper(new ObjectMapper()) + .messageEndpoint(MESSAGE_ENDPOINT) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .build(); } @Bean diff --git a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseSyncServerTransportTests.java b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseSyncServerTransportTests.java index 7e49ddf3b..101a067ad 100644 --- a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseSyncServerTransportTests.java +++ b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcSseSyncServerTransportTests.java @@ -36,7 +36,10 @@ static class TestConfig { @Bean public WebMvcSseServerTransportProvider webMvcSseServerTransportProvider() { - return new WebMvcSseServerTransportProvider(new ObjectMapper(), MESSAGE_ENDPOINT); + return WebMvcSseServerTransportProvider.builder() + .objectMapper(new ObjectMapper()) + .messageEndpoint(MESSAGE_ENDPOINT) + .build(); } @Bean diff --git a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStatelessIntegrationTests.java b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStatelessIntegrationTests.java index b2264ea00..c7c1e710d 100644 --- a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStatelessIntegrationTests.java +++ b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStatelessIntegrationTests.java @@ -11,8 +11,7 @@ import org.apache.catalina.LifecycleState; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; +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,9 +28,9 @@ import io.modelcontextprotocol.server.McpServer.StatelessAsyncSpecification; import io.modelcontextprotocol.server.McpServer.StatelessSyncSpecification; import io.modelcontextprotocol.server.transport.WebMvcStatelessServerTransport; -import io.modelcontextprotocol.spec.McpSchema; import reactor.core.scheduler.Schedulers; +@Timeout(15) class WebMvcStatelessIntegrationTests extends AbstractStatelessIntegrationTests { private static final int PORT = TestUtil.findAvailablePort(); @@ -63,6 +62,32 @@ public RouterFunction routerFunction(WebMvcStatelessServerTransp private TomcatTestUtil.TomcatServer tomcatServer; + @Override + protected StatelessAsyncSpecification prepareAsyncServerBuilder() { + return McpServer.async(this.mcpServerTransport); + } + + @Override + protected StatelessSyncSpecification prepareSyncServerBuilder() { + return McpServer.sync(this.mcpServerTransport); + } + + @Override + protected void prepareClients(int port, String mcpEndpoint) { + + clientBuilders.put("httpclient", McpClient + .sync(HttpClientStreamableHttpTransport.builder("http://localhost:" + port).endpoint(mcpEndpoint).build()) + .requestTimeout(Duration.ofHours(10))); + + clientBuilders.put("webflux", + McpClient + .sync(WebClientStreamableHttpTransport + .builder(WebClient.builder().baseUrl("http://localhost:" + port)) + .endpoint(mcpEndpoint) + .build()) + .requestTimeout(Duration.ofHours(10))); + } + @BeforeEach public void before() { @@ -76,33 +101,13 @@ public void before() { throw new RuntimeException("Failed to start Tomcat", e); } - clientBuilders - .put("httpclient", - McpClient.sync(HttpClientStreamableHttpTransport.builder("http://localhost:" + PORT) - .endpoint(MESSAGE_ENDPOINT) - .build()).initializationTimeout(Duration.ofHours(10)).requestTimeout(Duration.ofHours(10))); - - clientBuilders.put("webflux", - McpClient.sync(WebClientStreamableHttpTransport - .builder(WebClient.builder().baseUrl("http://localhost:" + PORT)) - .endpoint(MESSAGE_ENDPOINT) - .build())); + prepareClients(PORT, MESSAGE_ENDPOINT); // Get the transport from Spring context this.mcpServerTransport = tomcatServer.appContext().getBean(WebMvcStatelessServerTransport.class); } - @Override - protected StatelessAsyncSpecification prepareAsyncServerBuilder() { - return McpServer.async(this.mcpServerTransport); - } - - @Override - protected StatelessSyncSpecification prepareSyncServerBuilder() { - return McpServer.sync(this.mcpServerTransport); - } - @AfterEach public void after() { reactor.netty.http.HttpResources.disposeLoopsAndConnections(); @@ -124,42 +129,4 @@ public void after() { } } - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void simple(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var server = McpServer.async(this.mcpServerTransport) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(1000)) - .build(); - - try ( - // Create client without sampling capabilities - var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample " + "client", "0.0.0")) - .requestTimeout(Duration.ofSeconds(1000)) - .build()) { - - assertThat(client.initialize()).isNotNull(); - - } - server.closeGracefully(); - } - - @Override - protected void prepareClients(int port, String mcpEndpoint) { - - clientBuilders.put("httpclient", McpClient - .sync(HttpClientStreamableHttpTransport.builder("http://localhost:" + port).endpoint(mcpEndpoint).build()) - .initializationTimeout(Duration.ofHours(10)) - .requestTimeout(Duration.ofHours(10))); - - clientBuilders.put("webflux", - McpClient.sync(WebClientStreamableHttpTransport - .builder(WebClient.builder().baseUrl("http://localhost:" + port)) - .endpoint(mcpEndpoint) - .build())); - } - } diff --git a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java index f99b016ff..2f4c651fd 100644 --- a/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java +++ b/mcp-spring/mcp-spring-webmvc/src/test/java/io/modelcontextprotocol/server/WebMvcStreamableIntegrationTests.java @@ -11,11 +11,11 @@ import org.apache.catalina.LifecycleState; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; +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; +import org.springframework.web.servlet.function.ServerRequest; import org.springframework.web.servlet.config.annotation.EnableWebMvc; import org.springframework.web.servlet.function.RouterFunction; import org.springframework.web.servlet.function.ServerResponse; @@ -29,9 +29,9 @@ import io.modelcontextprotocol.server.McpServer.AsyncSpecification; import io.modelcontextprotocol.server.McpServer.SyncSpecification; import io.modelcontextprotocol.server.transport.WebMvcStreamableServerTransportProvider; -import io.modelcontextprotocol.spec.McpSchema; import reactor.core.scheduler.Schedulers; +@Timeout(15) class WebMvcStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests { private static final int PORT = TestUtil.findAvailablePort(); @@ -40,6 +40,11 @@ class WebMvcStreamableIntegrationTests extends AbstractMcpClientServerIntegratio private WebMvcStreamableServerTransportProvider mcpServerTransportProvider; + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + @Configuration @EnableWebMvc static class TestConfig { @@ -48,6 +53,7 @@ static class TestConfig { public WebMvcStreamableServerTransportProvider webMvcStreamableServerTransportProvider() { return WebMvcStreamableServerTransportProvider.builder() .objectMapper(new ObjectMapper()) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .mcpEndpoint(MESSAGE_ENDPOINT) .build(); } @@ -124,42 +130,20 @@ public void after() { } } - @ParameterizedTest(name = "{0} : {displayName} ") - @ValueSource(strings = { "httpclient", "webflux" }) - void simple(String clientType) { - - var clientBuilder = clientBuilders.get(clientType); - - var server = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(1000)) - .build(); - - try ( - // Create client without sampling capabilities - var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample " + "client", "0.0.0")) - .requestTimeout(Duration.ofSeconds(1000)) - .build()) { - - assertThat(client.initialize()).isNotNull(); - - } - server.closeGracefully(); - } - @Override protected void prepareClients(int port, String mcpEndpoint) { clientBuilders.put("httpclient", McpClient .sync(HttpClientStreamableHttpTransport.builder("http://localhost:" + port).endpoint(mcpEndpoint).build()) - .initializationTimeout(Duration.ofHours(10)) .requestTimeout(Duration.ofHours(10))); clientBuilders.put("webflux", - McpClient.sync(WebClientStreamableHttpTransport - .builder(WebClient.builder().baseUrl("http://localhost:" + port)) - .endpoint(mcpEndpoint) - .build())); + McpClient + .sync(WebClientStreamableHttpTransport + .builder(WebClient.builder().baseUrl("http://localhost:" + port)) + .endpoint(mcpEndpoint) + .build()) + .requestTimeout(Duration.ofHours(10))); } } diff --git a/mcp-test/pom.xml b/mcp-test/pom.xml index cc34e96d4..6e661bb6a 100644 --- a/mcp-test/pom.xml +++ b/mcp-test/pom.xml @@ -6,7 +6,7 @@ io.modelcontextprotocol.sdk mcp-parent - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT mcp-test jar @@ -24,7 +24,7 @@ io.modelcontextprotocol.sdk mcp - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT diff --git a/mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java b/mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java index d3d4fc071..5246c1e2d 100644 --- a/mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java +++ b/mcp-test/src/main/java/io/modelcontextprotocol/AbstractMcpClientServerIntegrationTests.java @@ -1,12 +1,14 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol; import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatExceptionOfType; +import static org.assertj.core.api.Assertions.assertWith; import static org.awaitility.Awaitility.await; import static org.mockito.Mockito.mock; @@ -18,10 +20,14 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; import java.util.function.Function; +import java.util.stream.Collectors; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; @@ -30,20 +36,28 @@ import io.modelcontextprotocol.server.McpServer; import io.modelcontextprotocol.server.McpServerFeatures; import io.modelcontextprotocol.server.McpSyncServer; +import io.modelcontextprotocol.server.McpSyncServerExchange; +import io.modelcontextprotocol.server.McpTransportContext; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.CallToolResult; import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities; +import io.modelcontextprotocol.spec.McpSchema.CompleteRequest; +import io.modelcontextprotocol.spec.McpSchema.CompleteResult; import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest; import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult; import io.modelcontextprotocol.spec.McpSchema.ElicitRequest; import io.modelcontextprotocol.spec.McpSchema.ElicitResult; import io.modelcontextprotocol.spec.McpSchema.InitializeResult; import io.modelcontextprotocol.spec.McpSchema.ModelPreferences; +import io.modelcontextprotocol.spec.McpSchema.Prompt; +import io.modelcontextprotocol.spec.McpSchema.PromptArgument; +import io.modelcontextprotocol.spec.McpSchema.PromptReference; import io.modelcontextprotocol.spec.McpSchema.Role; import io.modelcontextprotocol.spec.McpSchema.Root; import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; import io.modelcontextprotocol.spec.McpSchema.Tool; +import io.modelcontextprotocol.util.Utils; import net.javacrumbs.jsonunit.core.Option; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; @@ -134,6 +148,8 @@ void testCreateMessageSuccess(String clientType) { CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); + AtomicReference samplingResult = new AtomicReference<>(); + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) .callHandler((exchange, request) -> { @@ -149,37 +165,35 @@ void testCreateMessageSuccess(String clientType) { .build()) .build(); - StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }).verifyComplete(); - - return Mono.just(callResponse); + return exchange.createMessage(createMessageRequest) + .doOnNext(samplingResult::set) + .thenReturn(callResponse); }) .build(); - //@formatter:off - var mcpServer = prepareAsyncServerBuilder() - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); + var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build(); - try ( - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build()) {//@formatter:on + try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) + .capabilities(ClientCapabilities.builder().sampling().build()) + .sampling(samplingHandler) + .build()) { InitializeResult initResult = mcpClient.initialize(); assertThat(initResult).isNotNull(); CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - assertThat(response).isNotNull().isEqualTo(callResponse); + assertThat(response).isNotNull(); + assertThat(response).isEqualTo(callResponse); + + assertWith(samplingResult.get(), result -> { + assertThat(result).isNotNull(); + assertThat(result.role()).isEqualTo(Role.USER); + assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); + assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); + assertThat(result.model()).isEqualTo("MockModelName"); + assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); + }); } mcpServer.close(); } @@ -215,6 +229,8 @@ void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws Interr CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); + AtomicReference samplingResult = new AtomicReference<>(); + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) .callHandler((exchange, request) -> { @@ -230,16 +246,9 @@ void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws Interr .build()) .build(); - StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }).verifyComplete(); - - return Mono.just(callResponse); + return exchange.createMessage(createMessageRequest) + .doOnNext(samplingResult::set) + .thenReturn(callResponse); }) .build(); @@ -256,6 +265,15 @@ void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws Interr assertThat(response).isNotNull(); assertThat(response).isEqualTo(callResponse); + assertWith(samplingResult.get(), result -> { + assertThat(result).isNotNull(); + assertThat(result.role()).isEqualTo(Role.USER); + assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); + assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); + assertThat(result.model()).isEqualTo("MockModelName"); + assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); + }); + mcpClient.close(); mcpServer.close(); } @@ -302,16 +320,7 @@ void testCreateMessageWithRequestTimeoutFail(String clientType) throws Interrupt .build()) .build(); - StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }).verifyComplete(); - - return Mono.just(callResponse); + return exchange.createMessage(createMessageRequest).thenReturn(callResponse); }) .build(); @@ -325,7 +334,7 @@ void testCreateMessageWithRequestTimeoutFail(String clientType) throws Interrupt assertThatExceptionOfType(McpError.class).isThrownBy(() -> { mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - }).withMessageContaining("Timeout"); + }).withMessageContaining("1000ms"); mcpClient.close(); mcpServer.close(); @@ -342,19 +351,14 @@ void testCreateElicitationWithoutElicitationCapabilities(String clientType) { McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) - .callHandler((exchange, request) -> { - - exchange.createElicitation(mock(McpSchema.ElicitRequest.class)).block(); - - return Mono.just(mock(CallToolResult.class)); - }) + .callHandler((exchange, request) -> exchange.createElicitation(mock(ElicitRequest.class)) + .then(Mono.just(mock(CallToolResult.class)))) .build(); var server = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build(); - try ( - // Create client without elicitation capabilities - var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) { + // Create client without elicitation capabilities + try (var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) { assertThat(client.initialize()).isNotNull(); @@ -430,17 +434,10 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { var clientBuilder = clientBuilders.get(clientType); - Function elicitationHandler = request -> { + Function elicitationHandler = request -> { assertThat(request.message()).isNotEmpty(); assertThat(request.requestedSchema()).isNotNull(); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new McpSchema.ElicitResult(McpSchema.ElicitResult.Action.ACCEPT, - Map.of("message", request.message())); + return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); }; var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) @@ -451,6 +448,8 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); + AtomicReference resultRef = new AtomicReference<>(); + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) .callHandler((exchange, request) -> { @@ -461,13 +460,9 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) .build(); - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); + return exchange.createElicitation(elicitationRequest) + .doOnNext(resultRef::set) + .then(Mono.just(callResponse)); }) .build(); @@ -483,6 +478,11 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { assertThat(response).isNotNull(); assertThat(response).isEqualTo(callResponse); + assertWith(resultRef.get(), result -> { + assertThat(result).isNotNull(); + assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT); + assertThat(result.content().get("message")).isEqualTo("Test message"); + }); mcpClient.closeGracefully(); mcpServer.closeGracefully().block(); @@ -739,7 +739,6 @@ void testRootsServerCloseWithActiveSubscription(String clientType) { // --------------------------------------- // Tools Tests // --------------------------------------- - String emptyJsonSchema = """ { "$schema": "http://json-schema.org/draft-07/schema#", @@ -754,6 +753,7 @@ void testToolCallSuccess(String clientType) { var clientBuilder = clientBuilders.get(clientType); + var responseBodyIsNullOrBlank = new AtomicBoolean(false); var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) @@ -767,7 +767,7 @@ void testToolCallSuccess(String clientType) { .GET() .build(), HttpResponse.BodyHandlers.ofString()); String responseBody = response.body(); - assertThat(responseBody).isNotBlank(); + responseBodyIsNullOrBlank.set(!Utils.hasText(responseBody)); } catch (Exception e) { e.printStackTrace(); @@ -790,6 +790,7 @@ void testToolCallSuccess(String clientType) { CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); + assertThat(responseBodyIsNullOrBlank.get()).isFalse(); assertThat(response).isNotNull().isEqualTo(callResponse); } @@ -833,6 +834,62 @@ void testThrowingToolCallIsCaughtBeforeTimeout(String clientType) { mcpServer.close(); } + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testToolCallSuccessWithTranportContextExtraction(String clientType) { + + var clientBuilder = clientBuilders.get(clientType); + + var transportContextIsNull = new AtomicBoolean(false); + var transportContextIsEmpty = new AtomicBoolean(false); + var responseBodyIsNullOrBlank = new AtomicBoolean(false); + + var expectedCallResponse = new McpSchema.CallToolResult( + List.of(new McpSchema.TextContent("CALL RESPONSE; ctx=value")), null); + McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() + .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) + .callHandler((exchange, request) -> { + + McpTransportContext transportContext = exchange.transportContext(); + transportContextIsNull.set(transportContext == null); + transportContextIsEmpty.set(transportContext.equals(McpTransportContext.EMPTY)); + String ctxValue = (String) transportContext.get("important"); + + try { + String responseBody = "TOOL RESPONSE"; + responseBodyIsNullOrBlank.set(!Utils.hasText(responseBody)); + } + catch (Exception e) { + e.printStackTrace(); + } + + return new McpSchema.CallToolResult( + List.of(new McpSchema.TextContent("CALL RESPONSE; ctx=" + ctxValue)), null); + }) + .build(); + + var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().tools(true).build()) + .tools(tool1) + .build(); + + try (var mcpClient = clientBuilder.build()) { + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); + + CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); + + assertThat(transportContextIsNull.get()).isFalse(); + assertThat(transportContextIsEmpty.get()).isFalse(); + assertThat(responseBodyIsNullOrBlank.get()).isFalse(); + assertThat(response).isNotNull().isEqualTo(expectedCallResponse); + } + + mcpServer.close(); + } + @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "httpclient", "webflux" }) void testToolListChangeHandlingSuccess(String clientType) { @@ -861,7 +918,7 @@ void testToolListChangeHandlingSuccess(String clientType) { }) .build(); - AtomicReference> rootsRef = new AtomicReference<>(); + AtomicReference> toolsRef = new AtomicReference<>(); var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().tools(true).build()) .tools(tool1) @@ -878,32 +935,31 @@ void testToolListChangeHandlingSuccess(String clientType) { .build(), HttpResponse.BodyHandlers.ofString()); String responseBody = response.body(); assertThat(responseBody).isNotBlank(); + toolsRef.set(toolsUpdate); } catch (Exception e) { e.printStackTrace(); } - - rootsRef.set(toolsUpdate); }).build()) { InitializeResult initResult = mcpClient.initialize(); assertThat(initResult).isNotNull(); - assertThat(rootsRef.get()).isNull(); + assertThat(toolsRef.get()).isNull(); assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); mcpServer.notifyToolsListChanged(); await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool1.tool())); + assertThat(toolsRef.get()).containsAll(List.of(tool1.tool())); }); // Remove a tool mcpServer.removeTool("tool1"); await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).isEmpty(); + assertThat(toolsRef.get()).isEmpty(); }); // Add a new tool @@ -919,7 +975,7 @@ void testToolListChangeHandlingSuccess(String clientType) { mcpServer.addTool(tool2); await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool2.tool())); + assertThat(toolsRef.get()).containsAll(List.of(tool2.tool())); }); } @@ -943,6 +999,276 @@ void testInitialize(String clientType) { mcpServer.close(); } + // --------------------------------------- + // Logging Tests + // --------------------------------------- + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testLoggingNotification(String clientType) throws InterruptedException { + int expectedNotificationsCount = 3; + CountDownLatch latch = new CountDownLatch(expectedNotificationsCount); + // Create a list to store received logging notifications + List receivedNotifications = new CopyOnWriteArrayList<>(); + + var clientBuilder = clientBuilders.get(clientType); + + // Create server with a tool that sends logging notifications + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() + .tool(Tool.builder() + .name("logging-test") + .description("Test logging notifications") + .inputSchema(emptyJsonSchema) + .build()) + .callHandler((exchange, request) -> { + + // Create and send notifications with different levels + + //@formatter:off + return exchange // This should be filtered out (DEBUG < NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.DEBUG) + .logger("test-logger") + .data("Debug message") + .build()) + .then(exchange // This should be sent (NOTICE >= NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.NOTICE) + .logger("test-logger") + .data("Notice message") + .build())) + .then(exchange // This should be sent (ERROR > NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.ERROR) + .logger("test-logger") + .data("Error message") + .build())) + .then(exchange // This should be filtered out (INFO < NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.INFO) + .logger("test-logger") + .data("Another info message") + .build())) + .then(exchange // This should be sent (ERROR >= NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.ERROR) + .logger("test-logger") + .data("Another error message") + .build())) + .thenReturn(new CallToolResult("Logging test completed", false)); + //@formatter:on + }) + .build(); + + var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0") + .capabilities(ServerCapabilities.builder().tools(true).build()) + .tools(tool) + .build(); + + try ( + // Create client with logging notification handler + var mcpClient = clientBuilder.loggingConsumer(notification -> { + receivedNotifications.add(notification); + latch.countDown(); + }).build()) { + + // Initialize client + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + // Set minimum logging level to NOTICE + mcpClient.setLoggingLevel(McpSchema.LoggingLevel.NOTICE); + + // Call the tool that sends logging notifications + CallToolResult result = mcpClient.callTool(new McpSchema.CallToolRequest("logging-test", Map.of())); + assertThat(result).isNotNull(); + assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); + assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Logging test completed"); + + assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue(); + + // Should have received 3 notifications (1 NOTICE and 2 ERROR) + assertThat(receivedNotifications).hasSize(expectedNotificationsCount); + + Map notificationMap = receivedNotifications.stream() + .collect(Collectors.toMap(n -> n.data(), n -> n)); + + // First notification should be NOTICE level + assertThat(notificationMap.get("Notice message").level()).isEqualTo(McpSchema.LoggingLevel.NOTICE); + assertThat(notificationMap.get("Notice message").logger()).isEqualTo("test-logger"); + assertThat(notificationMap.get("Notice message").data()).isEqualTo("Notice message"); + + // Second notification should be ERROR level + assertThat(notificationMap.get("Error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); + assertThat(notificationMap.get("Error message").logger()).isEqualTo("test-logger"); + assertThat(notificationMap.get("Error message").data()).isEqualTo("Error message"); + + // Third notification should be ERROR level + assertThat(notificationMap.get("Another error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); + assertThat(notificationMap.get("Another error message").logger()).isEqualTo("test-logger"); + assertThat(notificationMap.get("Another error message").data()).isEqualTo("Another error message"); + } + mcpServer.close(); + } + + // --------------------------------------- + // Progress Tests + // --------------------------------------- + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient", "webflux" }) + void testProgressNotification(String clientType) throws InterruptedException { + int expectedNotificationsCount = 4; // 3 notifications + 1 for another progress + // token + CountDownLatch latch = new CountDownLatch(expectedNotificationsCount); + // Create a list to store received logging notifications + List receivedNotifications = new CopyOnWriteArrayList<>(); + + var clientBuilder = clientBuilders.get(clientType); + + // Create server with a tool that sends logging notifications + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() + .tool(McpSchema.Tool.builder() + .name("progress-test") + .description("Test progress notifications") + .inputSchema(emptyJsonSchema) + .build()) + .callHandler((exchange, request) -> { + + // Create and send notifications + var progressToken = (String) request.meta().get("progressToken"); + + return exchange + .progressNotification( + new McpSchema.ProgressNotification(progressToken, 0.0, 1.0, "Processing started")) + .then(exchange.progressNotification( + new McpSchema.ProgressNotification(progressToken, 0.5, 1.0, "Processing data"))) + .then(// Send a progress notification with another progress value + // should + exchange.progressNotification(new McpSchema.ProgressNotification("another-progress-token", + 0.0, 1.0, "Another processing started"))) + .then(exchange.progressNotification( + new McpSchema.ProgressNotification(progressToken, 1.0, 1.0, "Processing completed"))) + .thenReturn(new CallToolResult(("Progress test completed"), false)); + }) + .build(); + + var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0") + .capabilities(ServerCapabilities.builder().tools(true).build()) + .tools(tool) + .build(); + + try ( + // Create client with progress notification handler + var mcpClient = clientBuilder.progressConsumer(notification -> { + receivedNotifications.add(notification); + latch.countDown(); + }).build()) { + + // Initialize client + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + // Call the tool that sends progress notifications + McpSchema.CallToolRequest callToolRequest = McpSchema.CallToolRequest.builder() + .name("progress-test") + .meta(Map.of("progressToken", "test-progress-token")) + .build(); + CallToolResult result = mcpClient.callTool(callToolRequest); + assertThat(result).isNotNull(); + assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); + assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Progress test completed"); + + assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue(); + + // Should have received 3 notifications + assertThat(receivedNotifications).hasSize(expectedNotificationsCount); + + Map notificationMap = receivedNotifications.stream() + .collect(Collectors.toMap(n -> n.message(), n -> n)); + + // First notification should be 0.0/1.0 progress + assertThat(notificationMap.get("Processing started").progressToken()).isEqualTo("test-progress-token"); + assertThat(notificationMap.get("Processing started").progress()).isEqualTo(0.0); + assertThat(notificationMap.get("Processing started").total()).isEqualTo(1.0); + assertThat(notificationMap.get("Processing started").message()).isEqualTo("Processing started"); + + // Second notification should be 0.5/1.0 progress + assertThat(notificationMap.get("Processing data").progressToken()).isEqualTo("test-progress-token"); + assertThat(notificationMap.get("Processing data").progress()).isEqualTo(0.5); + assertThat(notificationMap.get("Processing data").total()).isEqualTo(1.0); + assertThat(notificationMap.get("Processing data").message()).isEqualTo("Processing data"); + + // Third notification should be another progress token with 0.0/1.0 progress + assertThat(notificationMap.get("Another processing started").progressToken()) + .isEqualTo("another-progress-token"); + assertThat(notificationMap.get("Another processing started").progress()).isEqualTo(0.0); + assertThat(notificationMap.get("Another processing started").total()).isEqualTo(1.0); + assertThat(notificationMap.get("Another processing started").message()) + .isEqualTo("Another processing started"); + + // Fourth notification should be 1.0/1.0 progress + assertThat(notificationMap.get("Processing completed").progressToken()).isEqualTo("test-progress-token"); + assertThat(notificationMap.get("Processing completed").progress()).isEqualTo(1.0); + assertThat(notificationMap.get("Processing completed").total()).isEqualTo(1.0); + assertThat(notificationMap.get("Processing completed").message()).isEqualTo("Processing completed"); + } + finally { + mcpServer.close(); + } + } + + // --------------------------------------- + // Completion Tests + // --------------------------------------- + @ParameterizedTest(name = "{0} : Completion call") + @ValueSource(strings = { "httpclient", "webflux" }) + void testCompletionShouldReturnExpectedSuggestions(String clientType) { + var clientBuilder = clientBuilders.get(clientType); + + var expectedValues = List.of("python", "pytorch", "pyside"); + var completionResponse = new McpSchema.CompleteResult(new CompleteResult.CompleteCompletion(expectedValues, 10, // total + true // hasMore + )); + + AtomicReference samplingRequest = new AtomicReference<>(); + BiFunction completionHandler = (mcpSyncServerExchange, + request) -> { + samplingRequest.set(request); + return completionResponse; + }; + + var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().completions().build()) + .prompts(new McpServerFeatures.SyncPromptSpecification( + new Prompt("code_review", "Code review", "this is code review prompt", + List.of(new PromptArgument("language", "Language", "string", false))), + (mcpSyncServerExchange, getPromptRequest) -> null)) + .completions(new McpServerFeatures.SyncCompletionSpecification( + new McpSchema.PromptReference("ref/prompt", "code_review", "Code review"), completionHandler)) + .build(); + + try (var mcpClient = clientBuilder.build()) { + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + CompleteRequest request = new CompleteRequest( + new PromptReference("ref/prompt", "code_review", "Code review"), + new CompleteRequest.CompleteArgument("language", "py")); + + CompleteResult result = mcpClient.completeCompletion(request); + + assertThat(result).isNotNull(); + + assertThat(samplingRequest.get().argument().name()).isEqualTo("language"); + assertThat(samplingRequest.get().argument().value()).isEqualTo("py"); + assertThat(samplingRequest.get().ref().type()).isEqualTo("ref/prompt"); + } + + mcpServer.close(); + } + + // --------------------------------------- + // Ping Tests + // --------------------------------------- @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "httpclient", "webflux" }) void testPingSuccess(String clientType) { @@ -1005,7 +1331,6 @@ void testPingSuccess(String clientType) { // --------------------------------------- // Tool Structured Output Schema Tests // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "httpclient", "webflux" }) void testStructuredOutputValidationSuccess(String clientType) { diff --git a/mcp-test/src/main/java/io/modelcontextprotocol/AbstractStatelessIntegrationTests.java b/mcp-test/src/main/java/io/modelcontextprotocol/AbstractStatelessIntegrationTests.java index a84d127aa..618247d61 100644 --- a/mcp-test/src/main/java/io/modelcontextprotocol/AbstractStatelessIntegrationTests.java +++ b/mcp-test/src/main/java/io/modelcontextprotocol/AbstractStatelessIntegrationTests.java @@ -1,6 +1,7 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol; import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; diff --git a/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java b/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java index 22e8f195b..ed34ebff6 100644 --- a/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java +++ b/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java @@ -1,6 +1,7 @@ /* * Copyright 2024-2024 the original author or authors. */ + package io.modelcontextprotocol.client; import eu.rekawek.toxiproxy.Proxy; diff --git a/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java b/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java index 067fbac2c..ea3739da5 100644 --- a/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java +++ b/mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java @@ -486,7 +486,8 @@ void testAddRoot() { void testAddRootWithNullValue() { withClient(createMcpTransport(), mcpAsyncClient -> { StepVerifier.create(mcpAsyncClient.addRoot(null)) - .consumeErrorWith(e -> assertThat(e).isInstanceOf(McpError.class).hasMessage("Root must not be null")) + .consumeErrorWith(e -> assertThat(e).isInstanceOf(IllegalArgumentException.class) + .hasMessage("Root must not be null")) .verify(); }); } @@ -505,7 +506,7 @@ void testRemoveRoot() { void testRemoveNonExistentRoot() { withClient(createMcpTransport(), mcpAsyncClient -> { StepVerifier.create(mcpAsyncClient.removeRoot("nonexistent-uri")) - .consumeErrorWith(e -> assertThat(e).isInstanceOf(McpError.class) + .consumeErrorWith(e -> assertThat(e).isInstanceOf(IllegalStateException.class) .hasMessage("Root with uri 'nonexistent-uri' not found")) .verify(); }); diff --git a/mcp-test/src/main/java/io/modelcontextprotocol/server/TestUtil.java b/mcp-test/src/main/java/io/modelcontextprotocol/server/TestUtil.java index 0085f31ed..dbbf1a537 100644 --- a/mcp-test/src/main/java/io/modelcontextprotocol/server/TestUtil.java +++ b/mcp-test/src/main/java/io/modelcontextprotocol/server/TestUtil.java @@ -1,6 +1,7 @@ /* * Copyright 2025 - 2025 the original author or authors. */ + package io.modelcontextprotocol.server; import java.io.IOException; diff --git a/mcp/pom.xml b/mcp/pom.xml index de4ee988a..389559408 100644 --- a/mcp/pom.xml +++ b/mcp/pom.xml @@ -6,7 +6,7 @@ io.modelcontextprotocol.sdk mcp-parent - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT mcp jar diff --git a/mcp/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java b/mcp/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java index e33fafa6a..2cc1c5dba 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java +++ b/mcp/src/main/java/io/modelcontextprotocol/client/LifecycleInitializer.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.client; import java.time.Duration; @@ -285,8 +289,7 @@ public Mono withIntitialization(String actionName, Function this.initializationRef.get()) .timeout(this.initializationTimeout) .onErrorResume(ex -> { - logger.warn("Failed to initialize", ex); - return Mono.error(new McpError("Client failed to initialize " + actionName)); + return Mono.error(new RuntimeException("Client failed to initialize " + actionName, ex)); }) .flatMap(operation); }); @@ -311,8 +314,10 @@ private Mono doInitialize(DefaultInitialization init initializeResult.instructions()); if (!this.protocolVersions.contains(initializeResult.protocolVersion())) { - return Mono.error(new McpError( - "Unsupported protocol version from the server: " + initializeResult.protocolVersion())); + return Mono.error(McpError.builder(-32602) + .message("Unsupported protocol version") + .data("Unsupported protocol version from the server: " + initializeResult.protocolVersion()) + .build()); } return mcpClientSession.sendNotification(McpSchema.METHOD_NOTIFICATION_INITIALIZED, null) diff --git a/mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java b/mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java index 73765122f..eb6d42f68 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java +++ b/mcp/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java @@ -1,6 +1,7 @@ /* * Copyright 2024-2024 the original author or authors. */ + package io.modelcontextprotocol.client; import java.time.Duration; @@ -189,7 +190,8 @@ public class McpAsyncClient { // Sampling Handler if (this.clientCapabilities.sampling() != null) { if (features.samplingHandler() == null) { - throw new McpError("Sampling handler must not be null when client capabilities include sampling"); + throw new IllegalArgumentException( + "Sampling handler must not be null when client capabilities include sampling"); } this.samplingHandler = features.samplingHandler(); requestHandlers.put(McpSchema.METHOD_SAMPLING_CREATE_MESSAGE, samplingCreateMessageHandler()); @@ -198,7 +200,8 @@ public class McpAsyncClient { // Elicitation Handler if (this.clientCapabilities.elicitation() != null) { if (features.elicitationHandler() == null) { - throw new McpError("Elicitation handler must not be null when client capabilities include elicitation"); + throw new IllegalArgumentException( + "Elicitation handler must not be null when client capabilities include elicitation"); } this.elicitationHandler = features.elicitationHandler(); requestHandlers.put(McpSchema.METHOD_ELICITATION_CREATE, elicitationCreateHandler()); @@ -271,9 +274,9 @@ public class McpAsyncClient { notificationHandlers.put(McpSchema.METHOD_NOTIFICATION_PROGRESS, asyncProgressNotificationHandler(progressConsumersFinal)); - this.initializer = new LifecycleInitializer(clientCapabilities, clientInfo, - List.of(transport.protocolVersion()), initializationTimeout, ctx -> new McpClientSession(requestTimeout, - transport, requestHandlers, notificationHandlers, con -> con.contextWrite(ctx))); + this.initializer = new LifecycleInitializer(clientCapabilities, clientInfo, transport.protocolVersions(), + initializationTimeout, ctx -> new McpClientSession(requestTimeout, transport, requestHandlers, + notificationHandlers, con -> con.contextWrite(ctx))); this.transport.setExceptionHandler(this.initializer::handleException); } @@ -412,15 +415,15 @@ public Mono ping() { public Mono addRoot(Root root) { if (root == null) { - return Mono.error(new McpError("Root must not be null")); + return Mono.error(new IllegalArgumentException("Root must not be null")); } if (this.clientCapabilities.roots() == null) { - return Mono.error(new McpError("Client must be configured with roots capabilities")); + return Mono.error(new IllegalStateException("Client must be configured with roots capabilities")); } if (this.roots.containsKey(root.uri())) { - return Mono.error(new McpError("Root with uri '" + root.uri() + "' already exists")); + return Mono.error(new IllegalStateException("Root with uri '" + root.uri() + "' already exists")); } this.roots.put(root.uri(), root); @@ -446,11 +449,11 @@ public Mono addRoot(Root root) { public Mono removeRoot(String rootUri) { if (rootUri == null) { - return Mono.error(new McpError("Root uri must not be null")); + return Mono.error(new IllegalArgumentException("Root uri must not be null")); } if (this.clientCapabilities.roots() == null) { - return Mono.error(new McpError("Client must be configured with roots capabilities")); + return Mono.error(new IllegalStateException("Client must be configured with roots capabilities")); } Root removed = this.roots.remove(rootUri); @@ -468,7 +471,7 @@ public Mono removeRoot(String rootUri) { } return Mono.empty(); } - return Mono.error(new McpError("Root with uri '" + rootUri + "' not found")); + return Mono.error(new IllegalStateException("Root with uri '" + rootUri + "' not found")); } /** @@ -539,7 +542,7 @@ private RequestHandler elicitationCreateHandler() { public Mono callTool(McpSchema.CallToolRequest callToolRequest) { return this.initializer.withIntitialization("calling tools", init -> { if (init.initializeResult().capabilities().tools() == null) { - return Mono.error(new McpError("Server does not provide tools capability")); + return Mono.error(new IllegalStateException("Server does not provide tools capability")); } return init.mcpSession() .sendRequest(McpSchema.METHOD_TOOLS_CALL, callToolRequest, CALL_TOOL_RESULT_TYPE_REF); @@ -568,7 +571,7 @@ public Mono listTools() { public Mono listTools(String cursor) { return this.initializer.withIntitialization("listing tools", init -> { if (init.initializeResult().capabilities().tools() == null) { - return Mono.error(new McpError("Server does not provide tools capability")); + return Mono.error(new IllegalStateException("Server does not provide tools capability")); } return init.mcpSession() .sendRequest(McpSchema.METHOD_TOOLS_LIST, new McpSchema.PaginatedRequest(cursor), @@ -632,7 +635,7 @@ public Mono listResources() { public Mono listResources(String cursor) { return this.initializer.withIntitialization("listing resources", init -> { if (init.initializeResult().capabilities().resources() == null) { - return Mono.error(new McpError("Server does not provide the resources capability")); + return Mono.error(new IllegalStateException("Server does not provide the resources capability")); } return init.mcpSession() .sendRequest(McpSchema.METHOD_RESOURCES_LIST, new McpSchema.PaginatedRequest(cursor), @@ -664,7 +667,7 @@ public Mono readResource(McpSchema.Resource resour public Mono readResource(McpSchema.ReadResourceRequest readResourceRequest) { return this.initializer.withIntitialization("reading resources", init -> { if (init.initializeResult().capabilities().resources() == null) { - return Mono.error(new McpError("Server does not provide the resources capability")); + return Mono.error(new IllegalStateException("Server does not provide the resources capability")); } return init.mcpSession() .sendRequest(McpSchema.METHOD_RESOURCES_READ, readResourceRequest, READ_RESOURCE_RESULT_TYPE_REF); @@ -702,7 +705,7 @@ public Mono listResourceTemplates() { public Mono listResourceTemplates(String cursor) { return this.initializer.withIntitialization("listing resource templates", init -> { if (init.initializeResult().capabilities().resources() == null) { - return Mono.error(new McpError("Server does not provide the resources capability")); + return Mono.error(new IllegalStateException("Server does not provide the resources capability")); } return init.mcpSession() .sendRequest(McpSchema.METHOD_RESOURCES_TEMPLATES_LIST, new McpSchema.PaginatedRequest(cursor), @@ -862,10 +865,13 @@ private NotificationHandler asyncLoggingNotificationHandler( */ public Mono setLoggingLevel(LoggingLevel loggingLevel) { if (loggingLevel == null) { - return Mono.error(new McpError("Logging level must not be null")); + return Mono.error(new IllegalArgumentException("Logging level must not be null")); } return this.initializer.withIntitialization("setting logging level", init -> { + if (init.initializeResult().capabilities().logging() == null) { + return Mono.error(new IllegalStateException("Server's Logging capabilities are not enabled!")); + } var params = new McpSchema.SetLevelRequest(loggingLevel); return init.mcpSession().sendRequest(McpSchema.METHOD_LOGGING_SET_LEVEL, params, OBJECT_TYPE_REF).then(); }); diff --git a/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java b/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java index 3fe88fec8..a92f26c4e 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java +++ b/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientSseClientTransport.java @@ -1,6 +1,7 @@ /* * Copyright 2024 - 2025 the original author or authors. */ + package io.modelcontextprotocol.client.transport; import java.io.IOException; @@ -9,6 +10,7 @@ import java.net.http.HttpRequest; 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; @@ -22,9 +24,10 @@ import io.modelcontextprotocol.client.transport.ResponseSubscribers.ResponseEvent; import io.modelcontextprotocol.spec.McpClientTransport; -import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage; +import io.modelcontextprotocol.spec.McpTransportException; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.Utils; import reactor.core.Disposable; @@ -61,7 +64,7 @@ */ public class HttpClientSseClientTransport implements McpClientTransport { - private static final String MCP_PROTOCOL_VERSION = "2024-11-05"; + private static final String MCP_PROTOCOL_VERSION = ProtocolVersions.MCP_2024_11_05; private static final String MCP_PROTOCOL_VERSION_HEADER_NAME = "MCP-Protocol-Version"; @@ -167,8 +170,7 @@ public HttpClientSseClientTransport(HttpClient.Builder clientBuilder, String bas @Deprecated(forRemoval = true) public HttpClientSseClientTransport(HttpClient.Builder clientBuilder, HttpRequest.Builder requestBuilder, String baseUri, String sseEndpoint, ObjectMapper objectMapper) { - this(clientBuilder.connectTimeout(Duration.ofSeconds(10)).build(), requestBuilder, baseUri, sseEndpoint, - objectMapper); + this(clientBuilder.build(), requestBuilder, baseUri, sseEndpoint, objectMapper); } /** @@ -216,8 +218,8 @@ public HttpClientSseClientTransport(HttpClient.Builder clientBuilder, HttpReques } @Override - public String protocolVersion() { - return MCP_PROTOCOL_VERSION; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05); } /** @@ -238,17 +240,16 @@ public static class Builder { private String sseEndpoint = DEFAULT_SSE_ENDPOINT; - private HttpClient.Builder clientBuilder = HttpClient.newBuilder() - .version(HttpClient.Version.HTTP_1_1) - .connectTimeout(Duration.ofSeconds(10)); + private HttpClient.Builder clientBuilder = HttpClient.newBuilder().version(HttpClient.Version.HTTP_1_1); private ObjectMapper objectMapper = new ObjectMapper(); - private HttpRequest.Builder requestBuilder = HttpRequest.newBuilder() - .header("Content-Type", "application/json"); + private HttpRequest.Builder requestBuilder = HttpRequest.newBuilder(); private AsyncHttpRequestCustomizer httpRequestCustomizer = AsyncHttpRequestCustomizer.NOOP; + private Duration connectTimeout = Duration.ofSeconds(10); + /** * Creates a new builder instance. */ @@ -380,13 +381,25 @@ public Builder asyncHttpRequestCustomizer(AsyncHttpRequestCustomizer asyncHttpRe return this; } + /** + * Sets the connection timeout for the HTTP client. + * @param connectTimeout the connection timeout duration + * @return this builder + */ + public Builder connectTimeout(Duration connectTimeout) { + Assert.notNull(connectTimeout, "connectTimeout must not be null"); + this.connectTimeout = connectTimeout; + return this; + } + /** * Builds a new {@link HttpClientSseClientTransport} instance. * @return a new transport instance */ public HttpClientSseClientTransport build() { - return new HttpClientSseClientTransport(clientBuilder.build(), requestBuilder, baseUri, sseEndpoint, - objectMapper, httpRequestCustomizer); + HttpClient httpClient = this.clientBuilder.connectTimeout(this.connectTimeout).build(); + return new HttpClientSseClientTransport(httpClient, requestBuilder, baseUri, sseEndpoint, objectMapper, + httpRequestCustomizer); } } @@ -428,7 +441,7 @@ public Mono connect(Function, Mono> h return Flux.empty(); // No further processing needed } else { - sink.error(new McpError("Failed to handle SSE endpoint event")); + sink.error(new RuntimeException("Failed to handle SSE endpoint event")); } } else if (MESSAGE_EVENT_TYPE.equals(responseEvent.sseEvent().event())) { @@ -443,8 +456,7 @@ else if (MESSAGE_EVENT_TYPE.equals(responseEvent.sseEvent().event())) { } } catch (IOException e) { - logger.error("Error processing SSE event", e); - sink.error(new McpError("Error processing SSE event")); + sink.error(new McpTransportException("Error processing SSE event", e)); } } return Flux.error( @@ -517,8 +529,7 @@ private Mono serializeMessage(final JSONRPCMessage message) { return Mono.just(objectMapper.writeValueAsString(message)); } catch (IOException e) { - // TODO: why McpError and not RuntimeException? - return Mono.error(new McpError("Failed to serialize message")); + return Mono.error(new McpTransportException("Failed to serialize message", e)); } }); } @@ -528,6 +539,7 @@ private Mono> sendHttpPost(final String endpoint, final Str return Mono.defer(() -> { var builder = this.requestBuilder.copy() .uri(requestUri) + .header("Content-Type", "application/json") .header(MCP_PROTOCOL_VERSION_HEADER_NAME, MCP_PROTOCOL_VERSION) .POST(HttpRequest.BodyPublishers.ofString(body)); return Mono.from(this.httpRequestCustomizer.customize(builder, "POST", requestUri, body)); diff --git a/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java b/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java index dadb09abc..d8c49ae2f 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java +++ b/mcp/src/main/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransport.java @@ -30,11 +30,12 @@ import io.modelcontextprotocol.spec.DefaultMcpTransportStream; 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; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.Utils; import reactor.core.Disposable; @@ -73,7 +74,7 @@ public class HttpClientStreamableHttpTransport implements McpClientTransport { private static final Logger logger = LoggerFactory.getLogger(HttpClientStreamableHttpTransport.class); - private static final String MCP_PROTOCOL_VERSION = "2025-03-26"; + private static final String MCP_PROTOCOL_VERSION = ProtocolVersions.MCP_2025_03_26; private static final String DEFAULT_ENDPOINT = "/mcp"; @@ -135,8 +136,8 @@ private HttpClientStreamableHttpTransport(ObjectMapper objectMapper, HttpClient } @Override - public String protocolVersion() { - return MCP_PROTOCOL_VERSION; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05, ProtocolVersions.MCP_2025_03_26); } public static Builder builder(String baseUri) { @@ -287,9 +288,8 @@ private Mono reconnect(McpTransportStream stream) { } catch (IOException ioException) { - return Flux.error( - new McpError("Error parsing JSON-RPC message: " - + responseEvent.sseEvent().data())); + return Flux.error(new McpTransportException( + "Error parsing JSON-RPC message: " + responseEvent, ioException)); } } else { @@ -303,19 +303,39 @@ else if (statusCode == METHOD_NOT_ALLOWED) { // NotAllowed return Flux.empty(); } else if (statusCode == NOT_FOUND) { - String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession); - McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException( - "Session not found for session ID: " + sessionIdRepresentation); - return Flux.error(exception); + + 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); + } + return Flux.error( + new McpTransportException("Server Not Found. Status code:" + statusCode + + ", response-event:" + responseEvent)); } else if (statusCode == BAD_REQUEST) { - String sessionIdRepresentation = sessionIdOrPlaceholder(transportSession); - McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException( - "Session not found for session ID: " + sessionIdRepresentation); - return Flux.error(exception); + 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); + } + return Flux.error( + new McpTransportException("Bad Request. Status code:" + statusCode + + ", response-event:" + responseEvent)); + } - return Flux.error(new McpError( + return Flux.error(new McpTransportException( "Received unrecognized SSE event type: " + responseEvent.sseEvent().event())); }).flatMap( @@ -399,7 +419,7 @@ public Mono sendMessage(McpSchema.JSONRPCMessage sentMessage) { .header("Cache-Control", "no-cache") .header(HttpHeaders.PROTOCOL_VERSION, MCP_PROTOCOL_VERSION) .POST(HttpRequest.BodyPublishers.ofString(jsonBody)); - return Mono.from(this.httpRequestCustomizer.customize(builder, "GET", uri, jsonBody)); + return Mono.from(this.httpRequestCustomizer.customize(builder, "POST", uri, jsonBody)); }).flatMapMany(requestBuilder -> Flux.create(responseEventSink -> { // Create the async request with proper body subscriber selection @@ -467,8 +487,8 @@ else if (contentType.contains(TEXT_EVENT_STREAM)) { return Flux.from(sessionStream.consumeSseStream(Flux.just(idWithMessages))); } catch (IOException ioException) { - return Flux.error( - new McpError("Error parsing JSON-RPC message: " + sseEvent.data())); + return Flux.error(new McpTransportException( + "Error parsing JSON-RPC message: " + responseEvent, ioException)); } }); } @@ -484,8 +504,8 @@ else if (contentType.contains(APPLICATION_JSON)) { return Mono.just(McpSchema.deserializeJsonRpcMessage(objectMapper, data)); } catch (IOException e) { - // TODO: this should be a McpTransportError - return Mono.error(e); + return Mono.error(new McpTransportException( + "Error deserializing JSON-RPC message: " + responseEvent, e)); } } logger.warn("Unknown media type {} returned for POST in session {}", contentType, @@ -495,18 +515,32 @@ else if (contentType.contains(APPLICATION_JSON)) { new RuntimeException("Unknown media type returned: " + contentType)); } else if (statusCode == NOT_FOUND) { - McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException( - "Session not found for session ID: " + sessionRepresentation); - return Flux.error(exception); + 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)); } - // 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 else if (statusCode == BAD_REQUEST) { - McpTransportSessionNotFoundException exception = new McpTransportSessionNotFoundException( - "Session not found for session ID: " + sessionRepresentation); - return Flux.error(exception); + // 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)); } return Flux.error( @@ -554,9 +588,7 @@ public static class Builder { private ObjectMapper objectMapper; - private HttpClient.Builder clientBuilder = HttpClient.newBuilder() - .version(HttpClient.Version.HTTP_1_1) - .connectTimeout(Duration.ofSeconds(10)); + private HttpClient.Builder clientBuilder = HttpClient.newBuilder().version(HttpClient.Version.HTTP_1_1); private String endpoint = DEFAULT_ENDPOINT; @@ -568,6 +600,8 @@ public static class Builder { private AsyncHttpRequestCustomizer httpRequestCustomizer = AsyncHttpRequestCustomizer.NOOP; + private Duration connectTimeout = Duration.ofSeconds(10); + /** * Creates a new builder with the specified base URI. * @param baseUri the base URI of the MCP server @@ -704,6 +738,17 @@ public Builder asyncHttpRequestCustomizer(AsyncHttpRequestCustomizer asyncHttpRe return this; } + /** + * Sets the connection timeout for the HTTP client. + * @param connectTimeout the connection timeout duration + * @return this builder + */ + public Builder connectTimeout(Duration connectTimeout) { + Assert.notNull(connectTimeout, "connectTimeout must not be null"); + this.connectTimeout = connectTimeout; + return this; + } + /** * Construct a fresh instance of {@link HttpClientStreamableHttpTransport} using * the current builder configuration. @@ -712,8 +757,10 @@ public Builder asyncHttpRequestCustomizer(AsyncHttpRequestCustomizer asyncHttpRe public HttpClientStreamableHttpTransport build() { ObjectMapper objectMapper = this.objectMapper != null ? this.objectMapper : new ObjectMapper(); - return new HttpClientStreamableHttpTransport(objectMapper, clientBuilder.build(), requestBuilder, baseUri, - endpoint, resumableStreams, openConnectionOnStartup, httpRequestCustomizer); + HttpClient httpClient = this.clientBuilder.connectTimeout(this.connectTimeout).build(); + + return new HttpClientStreamableHttpTransport(objectMapper, httpClient, requestBuilder, baseUri, endpoint, + resumableStreams, openConnectionOnStartup, httpRequestCustomizer); } } diff --git a/mcp/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java b/mcp/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java index eb9d3c65c..29dc23c35 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java +++ b/mcp/src/main/java/io/modelcontextprotocol/client/transport/ResponseSubscribers.java @@ -1,6 +1,7 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol.client.transport; import java.net.http.HttpResponse; @@ -11,8 +12,10 @@ import org.reactivestreams.FlowAdapters; import org.reactivestreams.Subscription; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; -import io.modelcontextprotocol.spec.McpError; +import io.modelcontextprotocol.spec.McpTransportException; import reactor.core.publisher.BaseSubscriber; import reactor.core.publisher.FluxSink; @@ -30,6 +33,8 @@ */ class ResponseSubscribers { + private static final Logger logger = LoggerFactory.getLogger(ResponseSubscribers.class); + record SseEvent(String id, String event, String data) { } @@ -136,7 +141,6 @@ protected void hookOnSubscribe(Subscription subscription) { @Override protected void hookOnNext(String line) { - if (line.isEmpty()) { // Empty line means end of event if (this.eventBuilder.length() > 0) { @@ -153,23 +157,31 @@ protected void hookOnNext(String line) { if (matcher.find()) { this.eventBuilder.append(matcher.group(1).trim()).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()); } + 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 - // TODO: This should be a McpTransportError - this.sink.error(new McpError( + this.sink.error(new McpTransportException( "Invalid SSE response. Status code: " + this.responseInfo.statusCode() + " Line: " + line)); } @@ -211,6 +223,8 @@ static class AggregateSubscriber extends BaseSubscriber { */ private ResponseInfo responseInfo; + volatile boolean hasRequestedDemand = false; + /** * Creates a new JsonLineSubscriber that will emit parsed JSON-RPC messages. * @param sink the {@link FluxSink} to emit parsed {@link ResponseEvent} objects @@ -224,7 +238,13 @@ public AggregateSubscriber(ResponseInfo responseInfo, FluxSink si @Override protected void hookOnSubscribe(Subscription subscription) { - sink.onRequest(subscription::request); + + sink.onRequest(n -> { + if (!hasRequestedDemand) { + subscription.request(Long.MAX_VALUE); + } + hasRequestedDemand = true; + }); // Register disposal callback to cancel subscription when Flux is disposed sink.onDispose(subscription::cancel); @@ -237,10 +257,12 @@ protected void hookOnNext(String line) { @Override protected void hookOnComplete() { - if (this.eventBuilder.length() > 0) { + + if (hasRequestedDemand) { String data = this.eventBuilder.toString(); this.sink.next(new AggregateResponseEvent(responseInfo, data)); } + this.sink.complete(); } @@ -260,6 +282,8 @@ static class BodilessResponseLineSubscriber extends BaseSubscriber { private final ResponseInfo responseInfo; + volatile boolean hasRequestedDemand = false; + public BodilessResponseLineSubscriber(ResponseInfo responseInfo, FluxSink sink) { this.sink = sink; this.responseInfo = responseInfo; @@ -269,7 +293,10 @@ public BodilessResponseLineSubscriber(ResponseInfo responseInfo, FluxSink { - subscription.request(n); + if (!hasRequestedDemand) { + subscription.request(Long.MAX_VALUE); + } + hasRequestedDemand = true; }); // Register disposal callback to cancel subscription when Flux is disposed @@ -280,11 +307,13 @@ protected void hookOnSubscribe(Subscription subscription) { @Override protected void hookOnComplete() { - // 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)); + if (hasRequestedDemand) { + // emit dummy event to be able to inspect the response info + // this is a shortcut allowing for a more streamlined processing using + // operator composition instead of having to deal with the + // CompletableFuture along the Subscriber for inspecting the result + this.sink.next(new DummyEvent(responseInfo)); + } this.sink.complete(); } diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/DefaultMcpStatelessServerHandler.java b/mcp/src/main/java/io/modelcontextprotocol/server/DefaultMcpStatelessServerHandler.java index 234a1d4a0..2df3514b6 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/DefaultMcpStatelessServerHandler.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/DefaultMcpStatelessServerHandler.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.server; import io.modelcontextprotocol.spec.McpError; @@ -31,9 +35,17 @@ public Mono handleRequest(McpTransportContext transpo } return requestHandler.handle(transportContext, request.params()) .map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), result, null)) - .onErrorResume(t -> Mono.just(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, - new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR, t.getMessage(), - null)))); + .onErrorResume(t -> { + McpSchema.JSONRPCResponse.JSONRPCError error; + if (t instanceof McpError mcpError && mcpError.getJsonRpcError() != null) { + error = mcpError.getJsonRpcError(); + } + else { + error = new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR, + t.getMessage(), null); + } + return Mono.just(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, error)); + }); } @Override diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/DefaultMcpTransportContext.java b/mcp/src/main/java/io/modelcontextprotocol/server/DefaultMcpTransportContext.java index 300bdf711..9e18e189d 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/DefaultMcpTransportContext.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/DefaultMcpTransportContext.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.server; import java.util.Map; diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java index 9605fb3f2..a51c2e36c 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java @@ -132,7 +132,7 @@ public class McpAsyncServer { this.mcpTransportProvider = mcpTransportProvider; this.objectMapper = objectMapper; this.serverInfo = features.serverInfo(); - this.serverCapabilities = features.serverCapabilities(); + this.serverCapabilities = features.serverCapabilities().mutate().logging().build(); this.instructions = features.instructions(); this.tools.addAll(withStructuredOutputHandling(jsonSchemaValidator, features.tools())); this.resources.putAll(features.resources()); @@ -145,7 +145,7 @@ public class McpAsyncServer { Map> requestHandlers = prepareRequestHandlers(); Map notificationHandlers = prepareNotificationHandlers(features); - this.protocolVersions = List.of(mcpTransportProvider.protocolVersion()); + this.protocolVersions = mcpTransportProvider.protocolVersions(); mcpTransportProvider.setSessionFactory(transport -> new McpServerSession(UUID.randomUUID().toString(), requestTimeout, transport, this::asyncInitializeRequestHandler, requestHandlers, notificationHandlers)); @@ -157,7 +157,7 @@ public class McpAsyncServer { this.mcpTransportProvider = mcpTransportProvider; this.objectMapper = objectMapper; this.serverInfo = features.serverInfo(); - this.serverCapabilities = features.serverCapabilities(); + this.serverCapabilities = features.serverCapabilities().mutate().logging().build(); this.instructions = features.instructions(); this.tools.addAll(withStructuredOutputHandling(jsonSchemaValidator, features.tools())); this.resources.putAll(features.resources()); @@ -170,7 +170,7 @@ public class McpAsyncServer { Map> requestHandlers = prepareRequestHandlers(); Map notificationHandlers = prepareNotificationHandlers(features); - this.protocolVersions = List.of(mcpTransportProvider.protocolVersion()); + this.protocolVersions = mcpTransportProvider.protocolVersions(); mcpTransportProvider.setSessionFactory(new DefaultMcpStreamableServerSessionFactory(requestTimeout, this::asyncInitializeRequestHandler, requestHandlers, notificationHandlers)); diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpInitRequestHandler.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpInitRequestHandler.java index 609744637..13ff45a54 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpInitRequestHandler.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpInitRequestHandler.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.server; import io.modelcontextprotocol.spec.McpSchema; diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessAsyncServer.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessAsyncServer.java index 565c53f13..41e0e9588 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessAsyncServer.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessAsyncServer.java @@ -118,7 +118,7 @@ public class McpStatelessAsyncServer { requestHandlers.put(McpSchema.METHOD_COMPLETION_COMPLETE, completionCompleteRequestHandler()); } - this.protocolVersions = List.of(mcpTransport.protocolVersion()); + this.protocolVersions = new ArrayList<>(mcpTransport.protocolVersions()); McpStatelessServerHandler handler = new DefaultMcpStatelessServerHandler(requestHandlers, Map.of()); mcpTransport.setMcpHandler(handler); diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessNotificationHandler.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessNotificationHandler.java index d9269a59b..6db79a62c 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessNotificationHandler.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessNotificationHandler.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.server; import reactor.core.publisher.Mono; diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessRequestHandler.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessRequestHandler.java index a6bf0d073..e5c9e7c09 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessRequestHandler.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessRequestHandler.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.server; import reactor.core.publisher.Mono; diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessServerFeatures.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessServerFeatures.java index 8be59a779..60c1dbb65 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessServerFeatures.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessServerFeatures.java @@ -67,10 +67,7 @@ record Async(McpSchema.Implementation serverInfo, McpSchema.ServerCapabilities s this.serverCapabilities = (serverCapabilities != null) ? serverCapabilities : new McpSchema.ServerCapabilities(null, // completions null, // experimental - new McpSchema.ServerCapabilities.LoggingCapabilities(), // Enable - // logging - // by - // default + null, // currently statless server doesn't support set logging !Utils.isEmpty(prompts) ? new McpSchema.ServerCapabilities.PromptCapabilities(false) : null, !Utils.isEmpty(resources) ? new McpSchema.ServerCapabilities.ResourceCapabilities(false, false) : null, diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessServerHandler.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessServerHandler.java index 80884435e..7c4e23cfc 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessServerHandler.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpStatelessServerHandler.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.server; import io.modelcontextprotocol.spec.McpSchema; diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpSyncServer.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpSyncServer.java index 38f5128e4..5adda1a74 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpSyncServer.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpSyncServer.java @@ -4,7 +4,6 @@ package io.modelcontextprotocol.server; -import io.modelcontextprotocol.server.McpServerFeatures.AsyncToolSpecification; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.LoggingMessageNotification; import io.modelcontextprotocol.util.Assert; diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpTransportContext.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpTransportContext.java index 3d51bb6e2..1cd540f72 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpTransportContext.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpTransportContext.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.server; import java.util.Collections; diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/McpTransportContextExtractor.java b/mcp/src/main/java/io/modelcontextprotocol/server/McpTransportContextExtractor.java index 472de8195..97fcecf0d 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/McpTransportContextExtractor.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/McpTransportContextExtractor.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.server; /** diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java b/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java index 24e749fc3..582120e3f 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProvider.java @@ -1,12 +1,14 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol.server.transport; import java.io.BufferedReader; import java.io.IOException; import java.io.PrintWriter; import java.time.Duration; +import java.util.List; import java.util.Map; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; @@ -14,11 +16,15 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; +import io.modelcontextprotocol.server.DefaultMcpTransportContext; +import io.modelcontextprotocol.server.McpTransportContext; +import io.modelcontextprotocol.server.McpTransportContextExtractor; import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpServerSession; import io.modelcontextprotocol.spec.McpServerTransport; import io.modelcontextprotocol.spec.McpServerTransportProvider; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.KeepAliveScheduler; import jakarta.servlet.AsyncContext; @@ -99,6 +105,8 @@ public class HttpServletSseServerTransportProvider extends HttpServlet implement /** Map of active client sessions, keyed by session ID */ private final Map sessions = new ConcurrentHashMap<>(); + private McpTransportContextExtractor contextExtractor; + /** Flag indicating if the transport is in the process of shutting down */ private final AtomicBoolean isClosing = new AtomicBoolean(false); @@ -141,7 +149,7 @@ public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String m @Deprecated public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, String sseEndpoint) { - this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, null); + this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, null, (serverRequest, context) -> context); } /** @@ -160,11 +168,38 @@ public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String b @Deprecated public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, String sseEndpoint, Duration keepAliveInterval) { + this(objectMapper, baseUrl, messageEndpoint, sseEndpoint, keepAliveInterval, + (serverRequest, context) -> context); + } + + /** + * Creates a new HttpServletSseServerTransportProvider instance with a custom SSE + * endpoint. + * @param objectMapper The JSON object mapper to use for message + * serialization/deserialization + * @param baseUrl The base URL for the server transport + * @param messageEndpoint The endpoint path where clients will send their messages + * @param sseEndpoint The endpoint path where clients will establish SSE connections + * @param keepAliveInterval The interval for keep-alive pings, or null to disable + * keep-alive functionality + * @param contextExtractor The extractor for transport context from the request. + * @deprecated Use the builder {@link #builder()} instead for better configuration + * options. + */ + private HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String baseUrl, String messageEndpoint, + String sseEndpoint, Duration keepAliveInterval, + McpTransportContextExtractor contextExtractor) { + + Assert.notNull(objectMapper, "ObjectMapper must not be null"); + Assert.notNull(messageEndpoint, "messageEndpoint must not be null"); + Assert.notNull(sseEndpoint, "sseEndpoint must not be null"); + Assert.notNull(contextExtractor, "Context extractor must not be null"); this.objectMapper = objectMapper; this.baseUrl = baseUrl; this.messageEndpoint = messageEndpoint; this.sseEndpoint = sseEndpoint; + this.contextExtractor = contextExtractor; if (keepAliveInterval != null) { @@ -179,8 +214,8 @@ public HttpServletSseServerTransportProvider(ObjectMapper objectMapper, String b } @Override - public String protocolVersion() { - return "2024-11-05"; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05); } /** @@ -336,10 +371,13 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response) body.append(line); } + final McpTransportContext transportContext = this.contextExtractor.extract(request, + new DefaultMcpTransportContext()); McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body.toString()); // Process the message through the session's handle method - session.handle(message).block(); // Block for Servlet compatibility + // Block for Servlet compatibility + session.handle(message).contextWrite(ctx -> ctx.put(McpTransportContext.KEY, transportContext)).block(); response.setStatus(HttpServletResponse.SC_OK); } @@ -531,6 +569,8 @@ public static class Builder { private String sseEndpoint = DEFAULT_SSE_ENDPOINT; + private McpTransportContextExtractor contextExtractor = (serverRequest, context) -> context; + private Duration keepAliveInterval; /** @@ -580,6 +620,19 @@ public Builder sseEndpoint(String sseEndpoint) { return this; } + /** + * Sets the context extractor for extracting transport context from the request. + * @param contextExtractor The context extractor to use. Must not be null. + * @return this builder instance + * @throws IllegalArgumentException if contextExtractor is null + */ + public HttpServletSseServerTransportProvider.Builder contextExtractor( + McpTransportContextExtractor contextExtractor) { + Assert.notNull(contextExtractor, "Context extractor must not be null"); + this.contextExtractor = contextExtractor; + return this; + } + /** * Sets the interval for keep-alive pings. *

@@ -606,7 +659,7 @@ public HttpServletSseServerTransportProvider build() { throw new IllegalStateException("MessageEndpoint must be set"); } return new HttpServletSseServerTransportProvider(objectMapper, baseUrl, messageEndpoint, sseEndpoint, - keepAliveInterval); + keepAliveInterval, contextExtractor); } } diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java b/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java index 6805bf194..8b95ec607 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/transport/HttpServletStreamableServerTransportProvider.java @@ -28,6 +28,7 @@ import io.modelcontextprotocol.spec.McpStreamableServerSession; import io.modelcontextprotocol.spec.McpStreamableServerTransport; import io.modelcontextprotocol.spec.McpStreamableServerTransportProvider; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.util.Assert; import io.modelcontextprotocol.util.KeepAliveScheduler; import jakarta.servlet.AsyncContext; @@ -155,8 +156,8 @@ private HttpServletStreamableServerTransportProvider(ObjectMapper objectMapper, } @Override - public String protocolVersion() { - return "2025-03-26"; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05, ProtocolVersions.MCP_2025_03_26); } @Override diff --git a/mcp/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java b/mcp/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java index d2943b31d..af602f610 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java +++ b/mcp/src/main/java/io/modelcontextprotocol/server/transport/StdioServerTransportProvider.java @@ -10,6 +10,7 @@ import java.io.InputStreamReader; import java.io.OutputStream; import java.nio.charset.StandardCharsets; +import java.util.List; import java.util.concurrent.Executors; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Function; @@ -22,6 +23,7 @@ import io.modelcontextprotocol.spec.McpServerSession; import io.modelcontextprotocol.spec.McpServerTransport; import io.modelcontextprotocol.spec.McpServerTransportProvider; +import io.modelcontextprotocol.spec.ProtocolVersions; import io.modelcontextprotocol.util.Assert; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -89,8 +91,8 @@ public StdioServerTransportProvider(ObjectMapper objectMapper, InputStream input } @Override - public String protocolVersion() { - return "2024-11-05"; + public List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05); } @Override diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultJsonSchemaValidator.java b/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultJsonSchemaValidator.java index cd8fc9659..f4bdc02eb 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultJsonSchemaValidator.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultJsonSchemaValidator.java @@ -1,6 +1,7 @@ /* * Copyright 2024-2024 the original author or authors. */ + package io.modelcontextprotocol.spec; import java.util.Map; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpStreamableServerSessionFactory.java b/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpStreamableServerSessionFactory.java index 8533e69cf..f497afd43 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpStreamableServerSessionFactory.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpStreamableServerSessionFactory.java @@ -1,8 +1,11 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import io.modelcontextprotocol.server.McpNotificationHandler; import io.modelcontextprotocol.server.McpRequestHandler; -import reactor.core.publisher.Mono; import java.time.Duration; import java.util.Map; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java index 56cdeaf7f..fdb7bfd89 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportSession.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import org.reactivestreams.Publisher; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportStream.java b/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportStream.java index eb2b7edeb..8d63fb50d 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportStream.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpTransportStream.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import org.reactivestreams.Publisher; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/HttpHeaders.java b/mcp/src/main/java/io/modelcontextprotocol/spec/HttpHeaders.java index c1c4c7a7d..65b80957c 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/HttpHeaders.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/HttpHeaders.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; /** @@ -15,7 +19,7 @@ public interface HttpHeaders { /** * Identifies events within an SSE Stream. */ - String LAST_EVENT_ID = "last-event-id"; + String LAST_EVENT_ID = "Last-Event-ID"; /** * Identifies the MCP protocol version. diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/JsonSchemaValidator.java b/mcp/src/main/java/io/modelcontextprotocol/spec/JsonSchemaValidator.java index c95e627a9..572d7c043 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/JsonSchemaValidator.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/JsonSchemaValidator.java @@ -1,6 +1,7 @@ /* * Copyright 2024-2024 the original author or authors. */ + package io.modelcontextprotocol.spec; import java.util.Map; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java index cc7d2abf8..f7db3d7aa 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientSession.java @@ -221,7 +221,7 @@ private Mono handleIncomingNotification(McpSchema.JSONRPCNotification noti return Mono.defer(() -> { var handler = notificationHandlers.get(notification.method()); if (handler == null) { - logger.error("No handler registered for notification method: {}", notification.method()); + logger.warn("No handler registered for notification method: {}", notification); return Mono.empty(); } return handler.handle(notification.params()); diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientTransport.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientTransport.java index 5c3b33131..22aec831b 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientTransport.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpClientTransport.java @@ -1,6 +1,7 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol.spec; import java.util.function.Consumer; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpError.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpError.java index 13e43240b..6172d8637 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpError.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpError.java @@ -1,9 +1,11 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol.spec; import io.modelcontextprotocol.spec.McpSchema.JSONRPCResponse.JSONRPCError; +import io.modelcontextprotocol.util.Assert; public class McpError extends RuntimeException { @@ -14,6 +16,7 @@ public McpError(JSONRPCError jsonRpcError) { this.jsonRpcError = jsonRpcError; } + @Deprecated public McpError(Object error) { super(error.toString()); } @@ -22,4 +25,37 @@ public JSONRPCError getJsonRpcError() { return jsonRpcError; } + public static Builder builder(int errorCode) { + return new Builder(errorCode); + } + + public static class Builder { + + private final int code; + + private String message; + + private Object data; + + private Builder(int code) { + this.code = code; + } + + public Builder message(String message) { + this.message = message; + return this; + } + + public Builder data(Object data) { + this.data = data; + return this; + } + + public McpError build() { + Assert.hasText(message, "message must not be empty"); + return new McpError(new JSONRPCError(code, message, data)); + } + + } + } \ No newline at end of file diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpLoggableSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpLoggableSession.java index ebc6e0949..f43a2c1d9 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpLoggableSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpLoggableSession.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; /** diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpSchema.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpSchema.java index fb4baabfb..3f8150271 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpSchema.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpSchema.java @@ -45,7 +45,7 @@ private McpSchema() { } @Deprecated - public static final String LATEST_PROTOCOL_VERSION = "2025-03-26"; + public static final String LATEST_PROTOCOL_VERSION = ProtocolVersions.MCP_2025_03_26; public static final String JSONRPC_VERSION = "2.0"; @@ -523,6 +523,21 @@ public record ResourceCapabilities(@JsonProperty("subscribe") Boolean subscribe, public record ToolCapabilities(@JsonProperty("listChanged") Boolean listChanged) { } + /** + * Create a mutated copy of this object with the specified changes. + * @return A new Builder instance with the same values as this object. + */ + public Builder mutate() { + var builder = new Builder(); + builder.completions = this.completions; + builder.experimental = this.experimental; + builder.logging = this.logging; + builder.prompts = this.prompts; + builder.resources = this.resources; + builder.tools = this.tools; + return builder; + } + public static Builder builder() { return new Builder(); } @@ -533,7 +548,7 @@ public static class Builder { private Map experimental; - private LoggingCapabilities logging = new LoggingCapabilities(); + private LoggingCapabilities logging; private PromptCapabilities prompts; @@ -2336,6 +2351,22 @@ public PromptReference(String name) { public String identifier() { return name(); } + + @Override + public boolean equals(Object obj) { + if (this == obj) + return true; + if (obj == null || getClass() != obj.getClass()) + return false; + PromptReference that = (PromptReference) obj; + return java.util.Objects.equals(identifier(), that.identifier()) + && java.util.Objects.equals(type(), that.type()); + } + + @Override + public int hashCode() { + return java.util.Objects.hash(identifier(), type()); + } } /** diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java index 0b0ef01cd..e562ca012 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import java.time.Duration; @@ -21,7 +25,7 @@ import reactor.core.publisher.Sinks; /** - * Represents a Model Control Protocol (MCP) session on the server side. It manages + * Represents a Model Context Protocol (MCP) session on the server side. It manages * bidirectional JSON-RPC communication with the client. */ public class McpServerSession implements McpLoggableSession { @@ -194,7 +198,9 @@ public Mono sendNotification(String method, Object params) { * @return a Mono that completes when the message is processed */ public Mono handle(McpSchema.JSONRPCMessage message) { - return Mono.defer(() -> { + return Mono.deferContextual(ctx -> { + McpTransportContext transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY); + // TODO handle errors for communication to without initialization happening // first if (message instanceof McpSchema.JSONRPCResponse response) { @@ -210,7 +216,7 @@ public Mono handle(McpSchema.JSONRPCMessage message) { } else if (message instanceof McpSchema.JSONRPCRequest request) { logger.debug("Received request: {}", request); - return handleIncomingRequest(request).onErrorResume(error -> { + return handleIncomingRequest(request, transportContext).onErrorResume(error -> { var errorResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), null, new McpSchema.JSONRPCResponse.JSONRPCError(McpSchema.ErrorCodes.INTERNAL_ERROR, error.getMessage(), null)); @@ -223,7 +229,7 @@ else if (message instanceof McpSchema.JSONRPCNotification notification) { // happening first logger.debug("Received notification: {}", notification); // TODO: in case of error, should the POST request be signalled? - return handleIncomingNotification(notification) + return handleIncomingNotification(notification, transportContext) .doOnError(error -> logger.error("Error handling notification: {}", error.getMessage())); } else { @@ -236,9 +242,11 @@ else if (message instanceof McpSchema.JSONRPCNotification notification) { /** * Handles an incoming JSON-RPC request by routing it to the appropriate handler. * @param request The incoming JSON-RPC request + * @param transportContext * @return A Mono containing the JSON-RPC response */ - private Mono handleIncomingRequest(McpSchema.JSONRPCRequest request) { + private Mono handleIncomingRequest(McpSchema.JSONRPCRequest request, + McpTransportContext transportContext) { return Mono.defer(() -> { Mono resultMono; if (McpSchema.METHOD_INITIALIZE.equals(request.method())) { @@ -262,7 +270,8 @@ private Mono handleIncomingRequest(McpSchema.JSONRPCR error.message(), error.data()))); } - resultMono = this.exchangeSink.asMono().flatMap(exchange -> handler.handle(exchange, request.params())); + resultMono = this.exchangeSink.asMono() + .flatMap(exchange -> handler.handle(copyExchange(exchange, transportContext), request.params())); } return resultMono .map(result -> new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), result, null)) @@ -276,27 +285,42 @@ private Mono handleIncomingRequest(McpSchema.JSONRPCR /** * Handles an incoming JSON-RPC notification by routing it to the appropriate handler. * @param notification The incoming JSON-RPC notification + * @param transportContext * @return A Mono that completes when the notification is processed */ - private Mono handleIncomingNotification(McpSchema.JSONRPCNotification notification) { + private Mono handleIncomingNotification(McpSchema.JSONRPCNotification notification, + McpTransportContext transportContext) { return Mono.defer(() -> { if (McpSchema.METHOD_NOTIFICATION_INITIALIZED.equals(notification.method())) { this.state.lazySet(STATE_INITIALIZED); // FIXME: The session ID passed here is not the same as the one in the // legacy SSE transport. exchangeSink.tryEmitValue(new McpAsyncServerExchange(this.id, this, clientCapabilities.get(), - clientInfo.get(), McpTransportContext.EMPTY)); + clientInfo.get(), transportContext)); } var handler = notificationHandlers.get(notification.method()); if (handler == null) { - logger.error("No handler registered for notification method: {}", notification.method()); + logger.warn("No handler registered for notification method: {}", notification); return Mono.empty(); } - return this.exchangeSink.asMono().flatMap(exchange -> handler.handle(exchange, notification.params())); + return this.exchangeSink.asMono() + .flatMap(exchange -> handler.handle(copyExchange(exchange, transportContext), notification.params())); }); } + /** + * This legacy implementation assumes an exchange is established upon the + * initialization phase see: exchangeSink.tryEmitValue(...), which creates a cached + * immutable exchange. Here, we create a new exchange and copy over everything from + * that cached exchange, and use it for a single HTTP request, with the transport + * context passed in. + */ + private McpAsyncServerExchange copyExchange(McpAsyncServerExchange exchange, McpTransportContext transportContext) { + return new McpAsyncServerExchange(exchange.sessionId(), this, exchange.getClientCapabilities(), + exchange.getClientInfo(), transportContext); + } + record MethodNotFoundError(String method, String message, Object data) { } diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransport.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransport.java index 632b8cee6..39c1644e0 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransport.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransport.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; /** diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransportProvider.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransportProvider.java index 382c0153b..02028ccdf 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransportProvider.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransportProvider.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; /** diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransportProviderBase.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransportProviderBase.java index 798575017..acb1ecac6 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransportProviderBase.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpServerTransportProviderBase.java @@ -1,5 +1,10 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; +import java.util.List; import java.util.Map; import reactor.core.publisher.Mono; @@ -59,8 +64,8 @@ default void close() { * Returns the protocol version supported by this transport provider. * @return the protocol version as a string */ - default String protocolVersion() { - return "2024-11-05"; + default List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05); } } diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpSession.java index 7b29ca651..3473a4da8 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpSession.java @@ -5,11 +5,10 @@ package io.modelcontextprotocol.spec; import com.fasterxml.jackson.core.type.TypeReference; -import io.modelcontextprotocol.util.Assert; import reactor.core.publisher.Mono; /** - * Represents a Model Control Protocol (MCP) session that handles communication between + * Represents a Model Context Protocol (MCP) session that handles communication between * clients and the server. This interface provides methods for sending requests and * notifications, as well as managing the session lifecycle. * diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStatelessServerTransport.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStatelessServerTransport.java index 329908469..c1234b130 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStatelessServerTransport.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStatelessServerTransport.java @@ -1,5 +1,11 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; +import java.util.List; + import io.modelcontextprotocol.server.McpStatelessServerHandler; import reactor.core.publisher.Mono; @@ -22,8 +28,8 @@ default void close() { */ Mono closeGracefully(); - default String protocolVersion() { - return "2025-03-26"; + default List protocolVersions() { + return List.of(ProtocolVersions.MCP_2025_03_26); } } diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java index c9b041fd6..ef7967c1e 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import java.time.Duration; @@ -193,7 +197,7 @@ public Mono accept(McpSchema.JSONRPCNotification notification) { McpTransportContext transportContext = ctx.getOrDefault(McpTransportContext.KEY, McpTransportContext.EMPTY); McpNotificationHandler notificationHandler = this.notificationHandlers.get(notification.method()); if (notificationHandler == null) { - logger.error("No handler registered for notification method: {}", notification.method()); + logger.warn("No handler registered for notification method: {}", notification); return Mono.empty(); } McpLoggableSession listeningStream = this.listeningStreamRef.get(); diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerTransport.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerTransport.java index 39e90ce86..f53c68900 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerTransport.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerTransport.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import reactor.core.publisher.Mono; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerTransportProvider.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerTransportProvider.java index b75081096..09fe9fb0e 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerTransportProvider.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerTransportProvider.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import reactor.core.publisher.Mono; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransport.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransport.java index 49c485059..1922548a6 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransport.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransport.java @@ -4,6 +4,8 @@ package io.modelcontextprotocol.spec; +import java.util.List; + import com.fasterxml.jackson.core.type.TypeReference; import io.modelcontextprotocol.spec.McpSchema.JSONRPCMessage; import reactor.core.publisher.Mono; @@ -77,8 +79,8 @@ default void close() { */ T unmarshalFrom(Object data, TypeReference typeRef); - default String protocolVersion() { - return "2024-11-05"; + default List protocolVersions() { + return List.of(ProtocolVersions.MCP_2024_11_05); } } diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportException.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportException.java new file mode 100644 index 000000000..cfd3dae31 --- /dev/null +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportException.java @@ -0,0 +1,38 @@ +/* +* Copyright 2025 - 2025 the original author or authors. +*/ +package io.modelcontextprotocol.spec; + +/** + * Exception thrown when there is an issue with the transport layer of the Model Context + * Protocol (MCP). + * + *

+ * This exception is used to indicate errors that occur during communication between the + * MCP client and server, such as connection failures, protocol violations, or unexpected + * responses. + * + * @author Christian Tzolov + */ +public class McpTransportException extends RuntimeException { + + private static final long serialVersionUID = 1L; + + public McpTransportException(String message) { + super(message); + } + + public McpTransportException(String message, Throwable cause) { + super(message, cause); + } + + public McpTransportException(Throwable cause) { + super(cause); + } + + public McpTransportException(String message, Throwable cause, boolean enableSuppression, + boolean writableStackTrace) { + super(message, cause, enableSuppression, writableStackTrace); + } + +} \ No newline at end of file diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportSession.java index 555f018f8..716ff0d16 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportSession.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import org.reactivestreams.Publisher; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportSessionNotFoundException.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportSessionNotFoundException.java index 474a18ae0..eced49ec3 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportSessionNotFoundException.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportSessionNotFoundException.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; /** diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportStream.java b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportStream.java index 2d6dcce75..322afda63 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportStream.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/McpTransportStream.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import org.reactivestreams.Publisher; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/MissingMcpTransportSession.java b/mcp/src/main/java/io/modelcontextprotocol/spec/MissingMcpTransportSession.java index c83f0bead..aa33a8167 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/spec/MissingMcpTransportSession.java +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/MissingMcpTransportSession.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; import com.fasterxml.jackson.core.type.TypeReference; diff --git a/mcp/src/main/java/io/modelcontextprotocol/spec/ProtocolVersions.java b/mcp/src/main/java/io/modelcontextprotocol/spec/ProtocolVersions.java new file mode 100644 index 000000000..d8cb913a5 --- /dev/null +++ b/mcp/src/main/java/io/modelcontextprotocol/spec/ProtocolVersions.java @@ -0,0 +1,23 @@ +package io.modelcontextprotocol.spec; + +public interface ProtocolVersions { + + /** + * MCP protocol version for 2024-11-05. + * https://modelcontextprotocol.io/specification/2024-11-05 + */ + String MCP_2024_11_05 = "2024-11-05"; + + /** + * MCP protocol version for 2025-03-26. + * https://modelcontextprotocol.io/specification/2025-03-26 + */ + String MCP_2025_03_26 = "2025-03-26"; + + /** + * MCP protocol version for 2025-06-18. + * https://modelcontextprotocol.io/specification/2025-06-18 + */ + String MCP_2025_06_18 = "2025-06-18"; + +} diff --git a/mcp/src/main/java/io/modelcontextprotocol/util/DeafaultMcpUriTemplateManagerFactory.java b/mcp/src/main/java/io/modelcontextprotocol/util/DeafaultMcpUriTemplateManagerFactory.java index 3870b76fc..44ea31690 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/util/DeafaultMcpUriTemplateManagerFactory.java +++ b/mcp/src/main/java/io/modelcontextprotocol/util/DeafaultMcpUriTemplateManagerFactory.java @@ -1,6 +1,7 @@ /* * Copyright 2025 - 2025 the original author or authors. */ + package io.modelcontextprotocol.util; /** diff --git a/mcp/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java b/mcp/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java index 30e8a2c2a..9d411cd41 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java +++ b/mcp/src/main/java/io/modelcontextprotocol/util/KeepAliveScheduler.java @@ -1,6 +1,7 @@ /** * Copyright 2025 - 2025 the original author or authors. */ + package io.modelcontextprotocol.util; import java.time.Duration; diff --git a/mcp/src/main/java/io/modelcontextprotocol/util/McpUriTemplateManagerFactory.java b/mcp/src/main/java/io/modelcontextprotocol/util/McpUriTemplateManagerFactory.java index 9644f9a6c..389727b45 100644 --- a/mcp/src/main/java/io/modelcontextprotocol/util/McpUriTemplateManagerFactory.java +++ b/mcp/src/main/java/io/modelcontextprotocol/util/McpUriTemplateManagerFactory.java @@ -1,6 +1,7 @@ /* * Copyright 2025 - 2025 the original author or authors. */ + package io.modelcontextprotocol.util; /** diff --git a/mcp/src/test/java/io/modelcontextprotocol/MockMcpClientTransport.java b/mcp/src/test/java/io/modelcontextprotocol/MockMcpClientTransport.java index b531d5739..b1113a6d0 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/MockMcpClientTransport.java +++ b/mcp/src/test/java/io/modelcontextprotocol/MockMcpClientTransport.java @@ -45,8 +45,8 @@ public MockMcpClientTransport withProtocolVersion(String protocolVersion) { } @Override - public String protocolVersion() { - return protocolVersion; + public List protocolVersions() { + return List.of(protocolVersion); } public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) { diff --git a/mcp/src/test/java/io/modelcontextprotocol/MockMcpServerTransportProvider.java b/mcp/src/test/java/io/modelcontextprotocol/MockMcpServerTransportProvider.java index 7ba35bbf0..e955be89f 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/MockMcpServerTransportProvider.java +++ b/mcp/src/test/java/io/modelcontextprotocol/MockMcpServerTransportProvider.java @@ -1,18 +1,7 @@ /* -* Copyright 2025 - 2025 the original author or authors. -* -* Licensed under the Apache License, Version 2.0 (the "License"); -* you may not use this file except in compliance with the License. -* You may obtain a copy of the License at -* -* https://www.apache.org/licenses/LICENSE-2.0 -* -* Unless required by applicable law or agreed to in writing, software -* distributed under the License is distributed on an "AS IS" BASIS, -* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -* See the License for the specific language governing permissions and -* limitations under the License. -*/ + * Copyright 2025-2025 the original author or authors. + */ + package io.modelcontextprotocol; import io.modelcontextprotocol.spec.McpSchema; diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java index b673ed612..ec23e21dc 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java @@ -1,6 +1,7 @@ /* * Copyright 2024-2024 the original author or authors. */ + package io.modelcontextprotocol.client; import eu.rekawek.toxiproxy.Proxy; diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java index e912e1dd6..3626d8ca0 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java @@ -487,7 +487,8 @@ void testAddRoot() { void testAddRootWithNullValue() { withClient(createMcpTransport(), mcpAsyncClient -> { StepVerifier.create(mcpAsyncClient.addRoot(null)) - .consumeErrorWith(e -> assertThat(e).isInstanceOf(McpError.class).hasMessage("Root must not be null")) + .consumeErrorWith(e -> assertThat(e).isInstanceOf(IllegalArgumentException.class) + .hasMessage("Root must not be null")) .verify(); }); } @@ -506,7 +507,7 @@ void testRemoveRoot() { void testRemoveNonExistentRoot() { withClient(createMcpTransport(), mcpAsyncClient -> { StepVerifier.create(mcpAsyncClient.removeRoot("nonexistent-uri")) - .consumeErrorWith(e -> assertThat(e).isInstanceOf(McpError.class) + .consumeErrorWith(e -> assertThat(e).isInstanceOf(IllegalStateException.class) .hasMessage("Root with uri 'nonexistent-uri' not found")) .verify(); }); diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpAsyncClientTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpAsyncClientTests.java index aa081b51b..aef2ab8dd 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpAsyncClientTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpAsyncClientTests.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.client; import org.junit.jupiter.api.Timeout; diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpSyncClientTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpSyncClientTests.java index 8285f417f..7f00de60e 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpSyncClientTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/HttpClientStreamableHttpSyncClientTests.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.client; import org.junit.jupiter.api.Timeout; diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java index c8d691924..02021edbf 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/LifecycleInitializerTests.java @@ -16,7 +16,6 @@ import org.mockito.MockitoAnnotations; import io.modelcontextprotocol.spec.McpClientSession; -import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException; import reactor.core.publisher.Mono; @@ -154,7 +153,7 @@ void shouldFailForUnsupportedProtocolVersion() { .thenReturn(Mono.just(unsupportedResult)); StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) - .expectError(McpError.class) + .expectError(RuntimeException.class) .verify(); verify(mockClientSession, never()).sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()); @@ -178,7 +177,7 @@ void shouldTimeoutOnSlowInitialization() { init -> Mono.just(init.initializeResult())), () -> virtualTimeScheduler, Long.MAX_VALUE) .expectSubscription() .expectNoEvent(INITIALIZE_TIMEOUT) - .expectError(McpError.class) + .expectError(RuntimeException.class) .verify(); } @@ -234,7 +233,7 @@ void shouldHandleInitializationFailure() { .thenReturn(Mono.error(new RuntimeException("Connection failed"))); StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))) - .expectError(McpError.class) + .expectError(RuntimeException.class) .verify(); assertThat(initializer.isInitialized()).isFalse(); diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientResponseHandlerTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientResponseHandlerTests.java index 11bd2e4e9..cab847512 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientResponseHandlerTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientResponseHandlerTests.java @@ -13,7 +13,6 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import io.modelcontextprotocol.MockMcpClientTransport; -import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities; import io.modelcontextprotocol.spec.McpSchema.InitializeResult; @@ -79,8 +78,9 @@ void testSuccessfulInitialization() { // Verify initialization result assertThat(result).isNotNull(); - assertThat(result.protocolVersion()).isEqualTo(transport.protocolVersion()); + assertThat(result.protocolVersion()).isEqualTo(transport.protocolVersions().get(0)); assertThat(result.capabilities()).isEqualTo(serverCapabilities); + assertThat(result.capabilities().logging()).isNull(); assertThat(result.serverInfo()).isEqualTo(serverInfo); assertThat(result.instructions()).isEqualTo("Test instructions"); @@ -373,7 +373,7 @@ void testSamplingCreateMessageRequestHandlingWithNullHandler() { // Create client with sampling capability but null handler assertThatThrownBy( () -> McpClient.async(transport).capabilities(ClientCapabilities.builder().sampling().build()).build()) - .isInstanceOf(McpError.class) + .isInstanceOf(IllegalArgumentException.class) .hasMessage("Sampling handler must not be null when client capabilities include sampling"); } @@ -521,7 +521,7 @@ void testElicitationCreateRequestHandlingWithNullHandler() { // Create client with elicitation capability but null handler assertThatThrownBy(() -> McpClient.async(transport) .capabilities(ClientCapabilities.builder().elicitation().build()) - .build()).isInstanceOf(McpError.class) + .build()).isInstanceOf(IllegalArgumentException.class) .hasMessage("Elicitation handler must not be null when client capabilities include elicitation"); } diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java index 7b6777cbe..ae33898b7 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java @@ -1,9 +1,15 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.client; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import io.modelcontextprotocol.spec.McpClientTransport; import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.ProtocolVersions; + import org.junit.jupiter.api.Test; import reactor.core.publisher.Mono; @@ -20,8 +26,8 @@ class McpAsyncClientTests { public static final McpSchema.ServerCapabilities MOCK_SERVER_CAPABILITIES = McpSchema.ServerCapabilities.builder() .build(); - public static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult("2024-11-05", - MOCK_SERVER_CAPABILITIES, MOCK_SERVER_INFO, "Test instructions"); + public static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult( + ProtocolVersions.MCP_2024_11_05, MOCK_SERVER_CAPABILITIES, MOCK_SERVER_INFO, "Test instructions"); private static final String CONTEXT_KEY = "context.key"; diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/McpClientProtocolVersionTests.java b/mcp/src/test/java/io/modelcontextprotocol/client/McpClientProtocolVersionTests.java index 2d41fc55f..3feb1d05c 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/McpClientProtocolVersionTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/McpClientProtocolVersionTests.java @@ -37,18 +37,20 @@ void shouldUseLatestVersionByDefault() { try { Mono initializeResultMono = client.initialize(); + String protocolVersion = transport.protocolVersions().get(transport.protocolVersions().size() - 1); + StepVerifier.create(initializeResultMono).then(() -> { McpSchema.JSONRPCRequest request = transport.getLastSentMessageAsRequest(); assertThat(request.params()).isInstanceOf(McpSchema.InitializeRequest.class); McpSchema.InitializeRequest initRequest = (McpSchema.InitializeRequest) request.params(); - assertThat(initRequest.protocolVersion()).isEqualTo(transport.protocolVersion()); + assertThat(initRequest.protocolVersion()).isEqualTo(transport.protocolVersions().get(0)); transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), - new McpSchema.InitializeResult(transport.protocolVersion(), null, + new McpSchema.InitializeResult(protocolVersion, null, new McpSchema.Implementation("test-server", "1.0.0"), null), null)); }).assertNext(result -> { - assertThat(result.protocolVersion()).isEqualTo(transport.protocolVersion()); + assertThat(result.protocolVersion()).isEqualTo(protocolVersion); }).verifyComplete(); } @@ -111,7 +113,7 @@ void shouldFailForUnsupportedVersion() { new McpSchema.InitializeResult(unsupportedVersion, null, new McpSchema.Implementation("test-server", "1.0.0"), null), null)); - }).expectError(McpError.class).verify(); + }).expectError(RuntimeException.class).verify(); } finally { StepVerifier.create(client.closeGracefully()).verifyComplete(); diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportEmptyJsonResponseTest.java b/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportEmptyJsonResponseTest.java new file mode 100644 index 000000000..8b3668671 --- /dev/null +++ b/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportEmptyJsonResponseTest.java @@ -0,0 +1,92 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + +package io.modelcontextprotocol.client.transport; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; + +import java.io.IOException; +import java.net.InetSocketAddress; +import java.net.URI; +import java.net.URISyntaxException; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import com.sun.net.httpserver.HttpServer; + +import io.modelcontextprotocol.server.transport.TomcatTestUtil; +import io.modelcontextprotocol.spec.McpSchema; +import io.modelcontextprotocol.spec.ProtocolVersions; +import reactor.test.StepVerifier; + +/** + * Handles emplty application/json response with 200 OK status code. + * + * @author codezkk + */ +public class HttpClientStreamableHttpTransportEmptyJsonResponseTest { + + static int PORT = TomcatTestUtil.findAvailablePort(); + + static String host = "http://localhost:" + PORT; + + static HttpServer server; + + @BeforeAll + static void startContainer() throws IOException { + + server = HttpServer.create(new InetSocketAddress(PORT), 0); + + // Empty, 200 OK response for the /mcp endpoint + server.createContext("/mcp", exchange -> { + exchange.getResponseHeaders().set("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, 0); + exchange.close(); + }); + + server.setExecutor(null); + server.start(); + } + + @AfterAll + static void stopContainer() { + server.stop(1); + } + + /** + * Regardless of the response (even if the response is null and the content-type is + * present), notify should handle it correctly. + */ + @Test + @Timeout(3) + void testNotificationInitialized() throws URISyntaxException { + + var uri = new URI(host + "/mcp"); + var mockRequestCustomizer = mock(SyncHttpRequestCustomizer.class); + var transport = HttpClientStreamableHttpTransport.builder(host) + .httpRequestCustomizer(mockRequestCustomizer) + .build(); + + var initializeRequest = new McpSchema.InitializeRequest(ProtocolVersions.MCP_2025_03_26, + McpSchema.ClientCapabilities.builder().roots(true).build(), + new McpSchema.Implementation("Spring AI MCP Client", "0.3.1")); + var testMessage = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, McpSchema.METHOD_INITIALIZE, + "test-id", initializeRequest); + + StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete(); + + // Verify the customizer was called + verify(mockRequestCustomizer, atLeastOnce()).customize(any(), eq("POST"), eq(uri), eq( + "{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":\"test-id\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{\"roots\":{\"listChanged\":true}},\"clientInfo\":{\"name\":\"Spring AI MCP Client\",\"version\":\"0.3.1\"}}}")); + + } + +} diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportErrorHandlingTest.java b/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportErrorHandlingTest.java new file mode 100644 index 000000000..2b502a83b --- /dev/null +++ b/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportErrorHandlingTest.java @@ -0,0 +1,345 @@ +/* + * 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.verify; + +import java.io.IOException; +import java.net.InetSocketAddress; +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.Test; +import org.junit.jupiter.api.Timeout; + +import com.sun.net.httpserver.HttpServer; + +import io.modelcontextprotocol.server.transport.TomcatTestUtil; +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.test.StepVerifier; + +/** + * Tests for error handling changes in HttpClientStreamableHttpTransport. Specifically + * tests the distinction between session-related errors and general transport errors for + * 404 and 400 status codes. + * + * @author Christian Tzolov + */ +@Timeout(15) +public class HttpClientStreamableHttpTransportErrorHandlingTest { + + private static final int PORT = TomcatTestUtil.findAvailablePort(); + + private static final String HOST = "http://localhost:" + PORT; + + private HttpServer server; + + private AtomicReference serverResponseStatus = new AtomicReference<>(200); + + private AtomicReference currentServerSessionId = new AtomicReference<>(null); + + private AtomicReference lastReceivedSessionId = new AtomicReference<>(null); + + private McpClientTransport transport; + + @BeforeEach + void startServer() throws IOException { + server = HttpServer.create(new InetSocketAddress(PORT), 0); + + // Configure the /mcp endpoint with dynamic response + server.createContext("/mcp", httpExchange -> { + if ("DELETE".equals(httpExchange.getRequestMethod())) { + httpExchange.sendResponseHeaders(200, 0); + } + else { + // Capture session ID from request if present + String requestSessionId = httpExchange.getRequestHeaders().getFirst(HttpHeaders.MCP_SESSION_ID); + lastReceivedSessionId.set(requestSessionId); + + int status = serverResponseStatus.get(); + + // Set response headers + httpExchange.getResponseHeaders().set("Content-Type", "application/json"); + + // Add session ID to response if configured + String responseSessionId = currentServerSessionId.get(); + if (responseSessionId != null) { + httpExchange.getResponseHeaders().set(HttpHeaders.MCP_SESSION_ID, responseSessionId); + } + + // Send response based on configured status + if (status == 200) { + String response = "{\"jsonrpc\":\"2.0\",\"result\":{},\"id\":\"test-id\"}"; + httpExchange.sendResponseHeaders(200, response.length()); + httpExchange.getResponseBody().write(response.getBytes()); + } + else { + httpExchange.sendResponseHeaders(status, 0); + } + } + httpExchange.close(); + }); + + server.setExecutor(null); + server.start(); + + transport = HttpClientStreamableHttpTransport.builder(HOST).build(); + } + + @AfterEach + void stopServer() { + if (server != null) { + server.stop(0); + } + } + + /** + * 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 = createTestRequestMessage(); + + StepVerifier.create(transport.sendMessage(testMessage)) + .expectErrorMatches(throwable -> throwable instanceof McpTransportException + && throwable.getMessage().contains("Not Found") && throwable.getMessage().contains("404") + && !(throwable instanceof McpTransportSessionNotFoundException)) + .verify(); + + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + } + + /** + * Test that 404 response WITH session ID throws McpTransportSessionNotFoundException + */ + @Test + void test404WithSessionId() { + // First establish a session + serverResponseStatus.set(200); + currentServerSessionId.set("test-session-123"); + + // Set up exception handler to verify session invalidation + @SuppressWarnings("unchecked") + Consumer 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 = createTestRequestMessage(); + StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete(); + + // The session should now be established, next request will include session ID + // Now return 404 for next request + serverResponseStatus.set(404); + + // Send another message - should get SessionNotFoundException + StepVerifier.create(transport.sendMessage(testMessage)) + .expectError(McpTransportSessionNotFoundException.class) + .verify(); + + // Verify exception handler was called with SessionNotFoundException + verify(exceptionHandler).accept(any(McpTransportSessionNotFoundException.class)); + + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + } + + /** + * 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 = createTestRequestMessage(); + + StepVerifier.create(transport.sendMessage(testMessage)) + .expectErrorMatches(throwable -> throwable instanceof McpTransportException + && throwable.getMessage().contains("Bad Request") && throwable.getMessage().contains("400") + && !(throwable instanceof McpTransportSessionNotFoundException)) + .verify(); + + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + } + + /** + * Test that 400 response WITH session ID throws McpTransportSessionNotFoundException + * This handles the case mentioned in the code comment about some implementations + * returning 400 for unknown session IDs. + */ + @Test + void test400WithSessionId() { + // First establish a session + serverResponseStatus.set(200); + currentServerSessionId.set("test-session-456"); + + // Set up exception handler + @SuppressWarnings("unchecked") + Consumer 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 = createTestRequestMessage(); + StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete(); + + // The session should now be established, next request will include session ID + // Now return 400 for next request (simulating unknown session ID) + serverResponseStatus.set(400); + + // Send another message - should get SessionNotFoundException + StepVerifier.create(transport.sendMessage(testMessage)) + .expectError(McpTransportSessionNotFoundException.class) + .verify(); + + // Verify exception handler was called + verify(exceptionHandler).accept(any(McpTransportSessionNotFoundException.class)); + + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + } + + /** + * Test session recovery after SessionNotFoundException Verifies that a new session + * can be established after the old one is invalidated + */ + @Test + void testSessionRecoveryAfter404() { + // First establish a session + serverResponseStatus.set(200); + currentServerSessionId.set("session-1"); + + // Send initial message to establish session + var testMessage = createTestRequestMessage(); + StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete(); + + assertThat(lastReceivedSessionId.get()).isNull(); + + // The session should now be established + // Simulate session loss - return 404 + serverResponseStatus.set(404); + + // This should fail with SessionNotFoundException + StepVerifier.create(transport.sendMessage(testMessage)) + .expectError(McpTransportSessionNotFoundException.class) + .verify(); + + // 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 + StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete(); + + // Verify no session ID was sent (since old session was invalidated) + assertThat(lastReceivedSessionId.get()).isNull(); + + // Next request should use the new session ID + StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete(); + + // Session ID should now be sent with requests + assertThat(lastReceivedSessionId.get()).isEqualTo("session-2"); + + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + } + + /** + * Test that reconnect (GET request) also properly handles 404/400 errors + */ + @Test + void testReconnectErrorHandling() { + + // 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)) { + 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 = HttpClientStreamableHttpTransport.builder(HOST) + .endpoint("/mcp-sse") + .openConnectionOnStartup(true) // This will trigger GET request on connect + .build(); + + // First connect successfully + StepVerifier.create(transport.connect(msg -> msg)).verifyComplete(); + + // Send message to establish session + var testMessage = createTestRequestMessage(); + StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete(); + + // Now simulate server returning 404 on reconnect + serverResponseStatus.set(404); + + // This should trigger reconnect which will fail + // The error should be handled internally and passed to exception handler + + StepVerifier.create(transport.closeGracefully()).verifyComplete(); + } + + private McpSchema.JSONRPCRequest createTestRequestMessage() { + 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); + } + +} diff --git a/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportTest.java b/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportTest.java index 479468f63..d645bb0b3 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportTest.java +++ b/mcp/src/test/java/io/modelcontextprotocol/client/transport/HttpClientStreamableHttpTransportTest.java @@ -80,7 +80,7 @@ void testRequestCustomizer() throws URISyntaxException { StepVerifier.create(t.sendMessage(testMessage)).verifyComplete(); // Verify the customizer was called - verify(mockRequestCustomizer, atLeastOnce()).customize(any(), eq("GET"), eq(uri), eq( + verify(mockRequestCustomizer, atLeastOnce()).customize(any(), eq("POST"), eq(uri), eq( "{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":\"test-id\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{\"roots\":{\"listChanged\":true}},\"clientInfo\":{\"name\":\"Spring AI MCP Client\",\"version\":\"0.3.1\"}}}")); }); } @@ -107,7 +107,7 @@ void testAsyncRequestCustomizer() throws URISyntaxException { StepVerifier.create(t.sendMessage(testMessage)).verifyComplete(); // Verify the customizer was called - verify(mockRequestCustomizer, atLeastOnce()).customize(any(), eq("GET"), eq(uri), eq( + verify(mockRequestCustomizer, atLeastOnce()).customize(any(), eq("POST"), eq(uri), eq( "{\"jsonrpc\":\"2.0\",\"method\":\"initialize\",\"id\":\"test-id\",\"params\":{\"protocolVersion\":\"2025-03-26\",\"capabilities\":{\"roots\":{\"listChanged\":true}},\"clientInfo\":{\"name\":\"Spring AI MCP Client\",\"version\":\"0.3.1\"}}}")); }); } diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java index 687ff6ae9..acaf0c8a9 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/AbstractMcpClientServerIntegrationTests.java @@ -1,12 +1,14 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol.server; import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatExceptionOfType; +import static org.assertj.core.api.Assertions.assertWith; import static org.awaitility.Awaitility.await; import static org.mockito.Mockito.mock; @@ -18,10 +20,14 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; import java.util.function.Function; +import java.util.stream.Collectors; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; @@ -31,16 +37,22 @@ import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.CallToolResult; import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities; +import io.modelcontextprotocol.spec.McpSchema.CompleteRequest; +import io.modelcontextprotocol.spec.McpSchema.CompleteResult; import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest; import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult; import io.modelcontextprotocol.spec.McpSchema.ElicitRequest; import io.modelcontextprotocol.spec.McpSchema.ElicitResult; import io.modelcontextprotocol.spec.McpSchema.InitializeResult; import io.modelcontextprotocol.spec.McpSchema.ModelPreferences; +import io.modelcontextprotocol.spec.McpSchema.Prompt; +import io.modelcontextprotocol.spec.McpSchema.PromptArgument; +import io.modelcontextprotocol.spec.McpSchema.PromptReference; import io.modelcontextprotocol.spec.McpSchema.Role; import io.modelcontextprotocol.spec.McpSchema.Root; import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; import io.modelcontextprotocol.spec.McpSchema.Tool; +import io.modelcontextprotocol.util.Utils; import net.javacrumbs.jsonunit.core.Option; import reactor.core.publisher.Mono; import reactor.test.StepVerifier; @@ -131,6 +143,8 @@ void testCreateMessageSuccess(String clientType) { CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); + AtomicReference samplingResult = new AtomicReference<>(); + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) .callHandler((exchange, request) -> { @@ -146,37 +160,35 @@ void testCreateMessageSuccess(String clientType) { .build()) .build(); - StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }).verifyComplete(); - - return Mono.just(callResponse); + return exchange.createMessage(createMessageRequest) + .doOnNext(samplingResult::set) + .thenReturn(callResponse); }) .build(); - //@formatter:off - var mcpServer = prepareAsyncServerBuilder() - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); + var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build(); - try ( - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build()) {//@formatter:on + try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) + .capabilities(ClientCapabilities.builder().sampling().build()) + .sampling(samplingHandler) + .build()) { InitializeResult initResult = mcpClient.initialize(); assertThat(initResult).isNotNull(); CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - assertThat(response).isNotNull().isEqualTo(callResponse); + assertThat(response).isNotNull(); + assertThat(response).isEqualTo(callResponse); + + assertWith(samplingResult.get(), result -> { + assertThat(result).isNotNull(); + assertThat(result.role()).isEqualTo(Role.USER); + assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); + assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); + assertThat(result.model()).isEqualTo("MockModelName"); + assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); + }); } mcpServer.close(); } @@ -212,6 +224,8 @@ void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws Interr CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); + AtomicReference samplingResult = new AtomicReference<>(); + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) .callHandler((exchange, request) -> { @@ -227,16 +241,9 @@ void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws Interr .build()) .build(); - StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }).verifyComplete(); - - return Mono.just(callResponse); + return exchange.createMessage(createMessageRequest) + .doOnNext(samplingResult::set) + .thenReturn(callResponse); }) .build(); @@ -253,6 +260,15 @@ void testCreateMessageWithRequestTimeoutSuccess(String clientType) throws Interr assertThat(response).isNotNull(); assertThat(response).isEqualTo(callResponse); + assertWith(samplingResult.get(), result -> { + assertThat(result).isNotNull(); + assertThat(result.role()).isEqualTo(Role.USER); + assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); + assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); + assertThat(result.model()).isEqualTo("MockModelName"); + assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); + }); + mcpClient.close(); mcpServer.close(); } @@ -299,16 +315,7 @@ void testCreateMessageWithRequestTimeoutFail(String clientType) throws Interrupt .build()) .build(); - StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }).verifyComplete(); - - return Mono.just(callResponse); + return exchange.createMessage(createMessageRequest).thenReturn(callResponse); }) .build(); @@ -322,7 +329,7 @@ void testCreateMessageWithRequestTimeoutFail(String clientType) throws Interrupt assertThatExceptionOfType(McpError.class).isThrownBy(() -> { mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - }).withMessageContaining("Timeout"); + }).withMessageContaining("1000ms"); mcpClient.close(); mcpServer.close(); @@ -339,19 +346,14 @@ void testCreateElicitationWithoutElicitationCapabilities(String clientType) { McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) - .callHandler((exchange, request) -> { - - exchange.createElicitation(mock(McpSchema.ElicitRequest.class)).block(); - - return Mono.just(mock(CallToolResult.class)); - }) + .callHandler((exchange, request) -> exchange.createElicitation(mock(ElicitRequest.class)) + .then(Mono.just(mock(CallToolResult.class)))) .build(); var server = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0").tools(tool).build(); - try ( - // Create client without elicitation capabilities - var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) { + // Create client without elicitation capabilities + try (var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) { assertThat(client.initialize()).isNotNull(); @@ -427,17 +429,10 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { var clientBuilder = clientBuilders.get(clientType); - Function elicitationHandler = request -> { + Function elicitationHandler = request -> { assertThat(request.message()).isNotEmpty(); assertThat(request.requestedSchema()).isNotNull(); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new McpSchema.ElicitResult(McpSchema.ElicitResult.Action.ACCEPT, - Map.of("message", request.message())); + return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); }; var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) @@ -448,6 +443,8 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); + AtomicReference resultRef = new AtomicReference<>(); + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) .callHandler((exchange, request) -> { @@ -458,13 +455,9 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) .build(); - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); + return exchange.createElicitation(elicitationRequest) + .doOnNext(resultRef::set) + .then(Mono.just(callResponse)); }) .build(); @@ -480,6 +473,11 @@ void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { assertThat(response).isNotNull(); assertThat(response).isEqualTo(callResponse); + assertWith(resultRef.get(), result -> { + assertThat(result).isNotNull(); + assertThat(result.action()).isEqualTo(McpSchema.ElicitResult.Action.ACCEPT); + assertThat(result.content().get("message")).isEqualTo("Test message"); + }); mcpClient.closeGracefully(); mcpServer.closeGracefully().block(); @@ -736,7 +734,6 @@ void testRootsServerCloseWithActiveSubscription(String clientType) { // --------------------------------------- // Tools Tests // --------------------------------------- - String emptyJsonSchema = """ { "$schema": "http://json-schema.org/draft-07/schema#", @@ -751,6 +748,7 @@ void testToolCallSuccess(String clientType) { var clientBuilder = clientBuilders.get(clientType); + var responseBodyIsNullOrBlank = new AtomicBoolean(false); var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) @@ -764,7 +762,7 @@ void testToolCallSuccess(String clientType) { .GET() .build(), HttpResponse.BodyHandlers.ofString()); String responseBody = response.body(); - assertThat(responseBody).isNotBlank(); + responseBodyIsNullOrBlank.set(!Utils.hasText(responseBody)); } catch (Exception e) { e.printStackTrace(); @@ -787,6 +785,7 @@ void testToolCallSuccess(String clientType) { CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); + assertThat(responseBodyIsNullOrBlank.get()).isFalse(); assertThat(response).isNotNull().isEqualTo(callResponse); } @@ -830,6 +829,62 @@ void testThrowingToolCallIsCaughtBeforeTimeout(String clientType) { mcpServer.close(); } + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient" }) + void testToolCallSuccessWithTranportContextExtraction(String clientType) { + + var clientBuilder = clientBuilders.get(clientType); + + var transportContextIsNull = new AtomicBoolean(false); + var transportContextIsEmpty = new AtomicBoolean(false); + var responseBodyIsNullOrBlank = new AtomicBoolean(false); + + var expectedCallResponse = new McpSchema.CallToolResult( + List.of(new McpSchema.TextContent("CALL RESPONSE; ctx=value")), null); + McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() + .tool(Tool.builder().name("tool1").description("tool1 description").inputSchema(emptyJsonSchema).build()) + .callHandler((exchange, request) -> { + + McpTransportContext transportContext = exchange.transportContext(); + transportContextIsNull.set(transportContext == null); + transportContextIsEmpty.set(transportContext.equals(McpTransportContext.EMPTY)); + String ctxValue = (String) transportContext.get("important"); + + try { + String responseBody = "TOOL RESPONSE"; + responseBodyIsNullOrBlank.set(!Utils.hasText(responseBody)); + } + catch (Exception e) { + e.printStackTrace(); + } + + return new McpSchema.CallToolResult( + List.of(new McpSchema.TextContent("CALL RESPONSE; ctx=" + ctxValue)), null); + }) + .build(); + + var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().tools(true).build()) + .tools(tool1) + .build(); + + try (var mcpClient = clientBuilder.build()) { + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); + + CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); + + assertThat(transportContextIsNull.get()).isFalse(); + assertThat(transportContextIsEmpty.get()).isFalse(); + assertThat(responseBodyIsNullOrBlank.get()).isFalse(); + assertThat(response).isNotNull().isEqualTo(expectedCallResponse); + } + + mcpServer.close(); + } + @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "httpclient" }) void testToolListChangeHandlingSuccess(String clientType) { @@ -858,7 +913,7 @@ void testToolListChangeHandlingSuccess(String clientType) { }) .build(); - AtomicReference> rootsRef = new AtomicReference<>(); + AtomicReference> toolsRef = new AtomicReference<>(); var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().tools(true).build()) .tools(tool1) @@ -875,32 +930,31 @@ void testToolListChangeHandlingSuccess(String clientType) { .build(), HttpResponse.BodyHandlers.ofString()); String responseBody = response.body(); assertThat(responseBody).isNotBlank(); + toolsRef.set(toolsUpdate); } catch (Exception e) { e.printStackTrace(); } - - rootsRef.set(toolsUpdate); }).build()) { InitializeResult initResult = mcpClient.initialize(); assertThat(initResult).isNotNull(); - assertThat(rootsRef.get()).isNull(); + assertThat(toolsRef.get()).isNull(); assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); mcpServer.notifyToolsListChanged(); await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool1.tool())); + assertThat(toolsRef.get()).containsAll(List.of(tool1.tool())); }); // Remove a tool mcpServer.removeTool("tool1"); await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).isEmpty(); + assertThat(toolsRef.get()).isEmpty(); }); // Add a new tool @@ -916,7 +970,7 @@ void testToolListChangeHandlingSuccess(String clientType) { mcpServer.addTool(tool2); await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool2.tool())); + assertThat(toolsRef.get()).containsAll(List.of(tool2.tool())); }); } @@ -940,6 +994,276 @@ void testInitialize(String clientType) { mcpServer.close(); } + // --------------------------------------- + // Logging Tests + // --------------------------------------- + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient" }) + void testLoggingNotification(String clientType) throws InterruptedException { + int expectedNotificationsCount = 3; + CountDownLatch latch = new CountDownLatch(expectedNotificationsCount); + // Create a list to store received logging notifications + List receivedNotifications = new CopyOnWriteArrayList<>(); + + var clientBuilder = clientBuilders.get(clientType); + + // Create server with a tool that sends logging notifications + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() + .tool(Tool.builder() + .name("logging-test") + .description("Test logging notifications") + .inputSchema(emptyJsonSchema) + .build()) + .callHandler((exchange, request) -> { + + // Create and send notifications with different levels + + //@formatter:off + return exchange // This should be filtered out (DEBUG < NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.DEBUG) + .logger("test-logger") + .data("Debug message") + .build()) + .then(exchange // This should be sent (NOTICE >= NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.NOTICE) + .logger("test-logger") + .data("Notice message") + .build())) + .then(exchange // This should be sent (ERROR > NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.ERROR) + .logger("test-logger") + .data("Error message") + .build())) + .then(exchange // This should be filtered out (INFO < NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.INFO) + .logger("test-logger") + .data("Another info message") + .build())) + .then(exchange // This should be sent (ERROR >= NOTICE) + .loggingNotification(McpSchema.LoggingMessageNotification.builder() + .level(McpSchema.LoggingLevel.ERROR) + .logger("test-logger") + .data("Another error message") + .build())) + .thenReturn(new CallToolResult("Logging test completed", false)); + //@formatter:on + }) + .build(); + + var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0") + .capabilities(ServerCapabilities.builder().logging().tools(true).build()) + .tools(tool) + .build(); + + try ( + // Create client with logging notification handler + var mcpClient = clientBuilder.loggingConsumer(notification -> { + receivedNotifications.add(notification); + latch.countDown(); + }).build()) { + + // Initialize client + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + // Set minimum logging level to NOTICE + mcpClient.setLoggingLevel(McpSchema.LoggingLevel.NOTICE); + + // Call the tool that sends logging notifications + CallToolResult result = mcpClient.callTool(new McpSchema.CallToolRequest("logging-test", Map.of())); + assertThat(result).isNotNull(); + assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); + assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Logging test completed"); + + assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue(); + + // Should have received 3 notifications (1 NOTICE and 2 ERROR) + assertThat(receivedNotifications).hasSize(expectedNotificationsCount); + + Map notificationMap = receivedNotifications.stream() + .collect(Collectors.toMap(n -> n.data(), n -> n)); + + // First notification should be NOTICE level + assertThat(notificationMap.get("Notice message").level()).isEqualTo(McpSchema.LoggingLevel.NOTICE); + assertThat(notificationMap.get("Notice message").logger()).isEqualTo("test-logger"); + assertThat(notificationMap.get("Notice message").data()).isEqualTo("Notice message"); + + // Second notification should be ERROR level + assertThat(notificationMap.get("Error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); + assertThat(notificationMap.get("Error message").logger()).isEqualTo("test-logger"); + assertThat(notificationMap.get("Error message").data()).isEqualTo("Error message"); + + // Third notification should be ERROR level + assertThat(notificationMap.get("Another error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); + assertThat(notificationMap.get("Another error message").logger()).isEqualTo("test-logger"); + assertThat(notificationMap.get("Another error message").data()).isEqualTo("Another error message"); + } + mcpServer.close(); + } + + // --------------------------------------- + // Progress Tests + // --------------------------------------- + @ParameterizedTest(name = "{0} : {displayName} ") + @ValueSource(strings = { "httpclient" }) + void testProgressNotification(String clientType) throws InterruptedException { + int expectedNotificationsCount = 4; // 3 notifications + 1 for another progress + // token + CountDownLatch latch = new CountDownLatch(expectedNotificationsCount); + // Create a list to store received logging notifications + List receivedNotifications = new CopyOnWriteArrayList<>(); + + var clientBuilder = clientBuilders.get(clientType); + + // Create server with a tool that sends logging notifications + McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() + .tool(McpSchema.Tool.builder() + .name("progress-test") + .description("Test progress notifications") + .inputSchema(emptyJsonSchema) + .build()) + .callHandler((exchange, request) -> { + + // Create and send notifications + var progressToken = (String) request.meta().get("progressToken"); + + return exchange + .progressNotification( + new McpSchema.ProgressNotification(progressToken, 0.0, 1.0, "Processing started")) + .then(exchange.progressNotification( + new McpSchema.ProgressNotification(progressToken, 0.5, 1.0, "Processing data"))) + .then(// Send a progress notification with another progress value + // should + exchange.progressNotification(new McpSchema.ProgressNotification("another-progress-token", + 0.0, 1.0, "Another processing started"))) + .then(exchange.progressNotification( + new McpSchema.ProgressNotification(progressToken, 1.0, 1.0, "Processing completed"))) + .thenReturn(new CallToolResult(("Progress test completed"), false)); + }) + .build(); + + var mcpServer = prepareAsyncServerBuilder().serverInfo("test-server", "1.0.0") + .capabilities(ServerCapabilities.builder().tools(true).build()) + .tools(tool) + .build(); + + try ( + // Create client with progress notification handler + var mcpClient = clientBuilder.progressConsumer(notification -> { + receivedNotifications.add(notification); + latch.countDown(); + }).build()) { + + // Initialize client + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + // Call the tool that sends progress notifications + McpSchema.CallToolRequest callToolRequest = McpSchema.CallToolRequest.builder() + .name("progress-test") + .meta(Map.of("progressToken", "test-progress-token")) + .build(); + CallToolResult result = mcpClient.callTool(callToolRequest); + assertThat(result).isNotNull(); + assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); + assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Progress test completed"); + + assertThat(latch.await(5, TimeUnit.SECONDS)).as("Should receive notifications in reasonable time").isTrue(); + + // Should have received 3 notifications + assertThat(receivedNotifications).hasSize(expectedNotificationsCount); + + Map notificationMap = receivedNotifications.stream() + .collect(Collectors.toMap(n -> n.message(), n -> n)); + + // First notification should be 0.0/1.0 progress + assertThat(notificationMap.get("Processing started").progressToken()).isEqualTo("test-progress-token"); + assertThat(notificationMap.get("Processing started").progress()).isEqualTo(0.0); + assertThat(notificationMap.get("Processing started").total()).isEqualTo(1.0); + assertThat(notificationMap.get("Processing started").message()).isEqualTo("Processing started"); + + // Second notification should be 0.5/1.0 progress + assertThat(notificationMap.get("Processing data").progressToken()).isEqualTo("test-progress-token"); + assertThat(notificationMap.get("Processing data").progress()).isEqualTo(0.5); + assertThat(notificationMap.get("Processing data").total()).isEqualTo(1.0); + assertThat(notificationMap.get("Processing data").message()).isEqualTo("Processing data"); + + // Third notification should be another progress token with 0.0/1.0 progress + assertThat(notificationMap.get("Another processing started").progressToken()) + .isEqualTo("another-progress-token"); + assertThat(notificationMap.get("Another processing started").progress()).isEqualTo(0.0); + assertThat(notificationMap.get("Another processing started").total()).isEqualTo(1.0); + assertThat(notificationMap.get("Another processing started").message()) + .isEqualTo("Another processing started"); + + // Fourth notification should be 1.0/1.0 progress + assertThat(notificationMap.get("Processing completed").progressToken()).isEqualTo("test-progress-token"); + assertThat(notificationMap.get("Processing completed").progress()).isEqualTo(1.0); + assertThat(notificationMap.get("Processing completed").total()).isEqualTo(1.0); + assertThat(notificationMap.get("Processing completed").message()).isEqualTo("Processing completed"); + } + finally { + mcpServer.close(); + } + } + + // --------------------------------------- + // Completion Tests + // --------------------------------------- + @ParameterizedTest(name = "{0} : Completion call") + @ValueSource(strings = { "httpclient" }) + void testCompletionShouldReturnExpectedSuggestions(String clientType) { + var clientBuilder = clientBuilders.get(clientType); + + var expectedValues = List.of("python", "pytorch", "pyside"); + var completionResponse = new McpSchema.CompleteResult(new CompleteResult.CompleteCompletion(expectedValues, 10, // total + true // hasMore + )); + + AtomicReference samplingRequest = new AtomicReference<>(); + BiFunction completionHandler = (mcpSyncServerExchange, + request) -> { + samplingRequest.set(request); + return completionResponse; + }; + + var mcpServer = prepareSyncServerBuilder().capabilities(ServerCapabilities.builder().completions().build()) + .prompts(new McpServerFeatures.SyncPromptSpecification( + new Prompt("code_review", "Code review", "this is code review prompt", + List.of(new PromptArgument("language", "Language", "string", false))), + (mcpSyncServerExchange, getPromptRequest) -> null)) + .completions(new McpServerFeatures.SyncCompletionSpecification( + new McpSchema.PromptReference("ref/prompt", "code_review", "Code review"), completionHandler)) + .build(); + + try (var mcpClient = clientBuilder.build()) { + + InitializeResult initResult = mcpClient.initialize(); + assertThat(initResult).isNotNull(); + + CompleteRequest request = new CompleteRequest( + new PromptReference("ref/prompt", "code_review", "Code review"), + new CompleteRequest.CompleteArgument("language", "py")); + + CompleteResult result = mcpClient.completeCompletion(request); + + assertThat(result).isNotNull(); + + assertThat(samplingRequest.get().argument().name()).isEqualTo("language"); + assertThat(samplingRequest.get().argument().value()).isEqualTo("py"); + assertThat(samplingRequest.get().ref().type()).isEqualTo("ref/prompt"); + } + + mcpServer.close(); + } + + // --------------------------------------- + // Ping Tests + // --------------------------------------- @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "httpclient" }) void testPingSuccess(String clientType) { @@ -1002,7 +1326,6 @@ void testPingSuccess(String clientType) { // --------------------------------------- // Tool Structured Output Schema Tests // --------------------------------------- - @ParameterizedTest(name = "{0} : {displayName} ") @ValueSource(strings = { "httpclient" }) void testStructuredOutputValidationSuccess(String clientType) { diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletSseIntegrationTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletSseIntegrationTests.java new file mode 100644 index 000000000..0f2991a9f --- /dev/null +++ b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletSseIntegrationTests.java @@ -0,0 +1,102 @@ +/* + * Copyright 2024 - 2024 the original author or authors. + */ + +package io.modelcontextprotocol.server; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.time.Duration; + +import org.apache.catalina.LifecycleException; +import org.apache.catalina.LifecycleState; +import org.apache.catalina.startup.Tomcat; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Timeout; + +import com.fasterxml.jackson.databind.ObjectMapper; + +import io.modelcontextprotocol.client.McpClient; +import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport; +import io.modelcontextprotocol.server.McpServer.AsyncSpecification; +import io.modelcontextprotocol.server.McpServer.SyncSpecification; +import io.modelcontextprotocol.server.transport.HttpServletSseServerTransportProvider; +import io.modelcontextprotocol.server.transport.TomcatTestUtil; +import jakarta.servlet.http.HttpServletRequest; + +@Timeout(15) +class HttpServletSseIntegrationTests extends AbstractMcpClientServerIntegrationTests { + + private static final int PORT = TomcatTestUtil.findAvailablePort(); + + private static final String CUSTOM_SSE_ENDPOINT = "/somePath/sse"; + + private static final String CUSTOM_MESSAGE_ENDPOINT = "/otherPath/mcp/message"; + + private HttpServletSseServerTransportProvider mcpServerTransportProvider; + + private Tomcat tomcat; + + @BeforeEach + public void before() { + // Create and configure the transport provider + mcpServerTransportProvider = HttpServletSseServerTransportProvider.builder() + .objectMapper(new ObjectMapper()) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) + .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) + .sseEndpoint(CUSTOM_SSE_ENDPOINT) + .build(); + + tomcat = TomcatTestUtil.createTomcatServer("", PORT, mcpServerTransportProvider); + try { + tomcat.start(); + assertThat(tomcat.getServer().getState()).isEqualTo(LifecycleState.STARTED); + } + catch (Exception e) { + throw new RuntimeException("Failed to start Tomcat", e); + } + + clientBuilders + .put("httpclient", + McpClient.sync(HttpClientSseClientTransport.builder("http://localhost:" + PORT) + .sseEndpoint(CUSTOM_SSE_ENDPOINT) + .build()).requestTimeout(Duration.ofHours(10))); + } + + @Override + protected AsyncSpecification prepareAsyncServerBuilder() { + return McpServer.async(this.mcpServerTransportProvider); + } + + @Override + protected SyncSpecification prepareSyncServerBuilder() { + return McpServer.sync(this.mcpServerTransportProvider); + } + + @AfterEach + public void after() { + if (mcpServerTransportProvider != null) { + mcpServerTransportProvider.closeGracefully().block(); + } + if (tomcat != null) { + try { + tomcat.stop(); + tomcat.destroy(); + } + catch (LifecycleException e) { + throw new RuntimeException("Failed to stop Tomcat", e); + } + } + } + + @Override + protected void prepareClients(int port, String mcpEndpoint) { + } + + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + +} diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStatelessIntegrationTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStatelessIntegrationTests.java index da8aa4adf..a8951e6dc 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStatelessIntegrationTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStatelessIntegrationTests.java @@ -1,35 +1,16 @@ /* * Copyright 2024 - 2024 the original author or authors. */ -package io.modelcontextprotocol.server; - -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json; -import static org.assertj.core.api.Assertions.assertThat; -import static org.awaitility.Awaitility.await; -import java.time.Duration; -import java.util.List; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.BiFunction; - -import org.apache.catalina.LifecycleException; -import org.apache.catalina.LifecycleState; -import org.apache.catalina.startup.Tomcat; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.params.ParameterizedTest; -import org.junit.jupiter.params.provider.ValueSource; -import org.springframework.web.client.RestClient; +package io.modelcontextprotocol.server; import com.fasterxml.jackson.databind.ObjectMapper; - import io.modelcontextprotocol.client.McpClient; import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport; import io.modelcontextprotocol.server.transport.HttpServletStatelessServerTransport; import io.modelcontextprotocol.server.transport.TomcatTestUtil; +import io.modelcontextprotocol.spec.HttpHeaders; +import io.modelcontextprotocol.spec.McpError; import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.CallToolResult; import io.modelcontextprotocol.spec.McpSchema.CompleteRequest; @@ -40,8 +21,36 @@ import io.modelcontextprotocol.spec.McpSchema.PromptReference; import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; import io.modelcontextprotocol.spec.McpSchema.Tool; +import io.modelcontextprotocol.spec.ProtocolVersions; import net.javacrumbs.jsonunit.core.Option; +import org.apache.catalina.LifecycleException; +import org.apache.catalina.LifecycleState; +import org.apache.catalina.startup.Tomcat; +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.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import org.springframework.mock.web.MockHttpServletRequest; +import org.springframework.mock.web.MockHttpServletResponse; +import org.springframework.web.client.RestClient; + +import java.time.Duration; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.BiFunction; +import static io.modelcontextprotocol.server.transport.HttpServletStatelessServerTransport.APPLICATION_JSON; +import static io.modelcontextprotocol.server.transport.HttpServletStatelessServerTransport.TEXT_EVENT_STREAM; +import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; +import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json; +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; + +@Timeout(15) class HttpServletStatelessIntegrationTests { private static final int PORT = TomcatTestUtil.findAvailablePort(); @@ -459,6 +468,49 @@ void testStructuredOutputRuntimeToolAddition(String clientType) { mcpServer.close(); } + @Test + void testThrownMcpError() throws Exception { + var mcpServer = McpServer.sync(mcpStatelessServerTransport) + .serverInfo("test-server", "1.0.0") + .capabilities(ServerCapabilities.builder().tools(true).build()) + .build(); + + Tool testTool = Tool.builder().name("test").description("test").build(); + + McpStatelessServerFeatures.SyncToolSpecification toolSpec = new McpStatelessServerFeatures.SyncToolSpecification( + testTool, (transportContext, request) -> { + throw new McpError(new McpSchema.JSONRPCResponse.JSONRPCError(12345, "testing", Map.of("a", "b"))); + }); + + mcpServer.addTool(toolSpec); + + McpSchema.CallToolRequest callToolRequest = new McpSchema.CallToolRequest("test", Map.of()); + McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, + McpSchema.METHOD_TOOLS_CALL, "test", callToolRequest); + + MockHttpServletRequest request = new MockHttpServletRequest("POST", CUSTOM_MESSAGE_ENDPOINT); + MockHttpServletResponse response = new MockHttpServletResponse(); + + byte[] content = new ObjectMapper().writeValueAsBytes(jsonrpcRequest); + request.setContent(content); + request.addHeader("Content-Type", "application/json"); + request.addHeader("Content-Length", Integer.toString(content.length)); + request.addHeader("Content-Length", Integer.toString(content.length)); + request.addHeader("Accept", APPLICATION_JSON + ", " + TEXT_EVENT_STREAM); + request.addHeader("Content-Type", APPLICATION_JSON); + request.addHeader("Cache-Control", "no-cache"); + request.addHeader(HttpHeaders.PROTOCOL_VERSION, ProtocolVersions.MCP_2025_03_26); + mcpStatelessServerTransport.service(request, response); + + McpSchema.JSONRPCResponse jsonrpcResponse = new ObjectMapper().readValue(response.getContentAsByteArray(), + McpSchema.JSONRPCResponse.class); + + assertThat(jsonrpcResponse.error()) + .isEqualTo(new McpSchema.JSONRPCResponse.JSONRPCError(12345, "testing", Map.of("a", "b"))); + + mcpServer.close(); + } + private double evaluateExpression(String expression) { // Simple expression evaluator for testing return switch (expression) { diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java index ecb0c33c3..2e9b4cbad 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/HttpServletStreamableIntegrationTests.java @@ -1,6 +1,7 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol.server; import static org.assertj.core.api.Assertions.assertThat; @@ -12,6 +13,7 @@ import org.apache.catalina.startup.Tomcat; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Timeout; import com.fasterxml.jackson.databind.ObjectMapper; @@ -21,7 +23,9 @@ import io.modelcontextprotocol.server.McpServer.SyncSpecification; import io.modelcontextprotocol.server.transport.HttpServletStreamableServerTransportProvider; import io.modelcontextprotocol.server.transport.TomcatTestUtil; +import jakarta.servlet.http.HttpServletRequest; +@Timeout(15) class HttpServletStreamableIntegrationTests extends AbstractMcpClientServerIntegrationTests { private static final int PORT = TomcatTestUtil.findAvailablePort(); @@ -37,6 +41,7 @@ public void before() { // Create and configure the transport provider mcpServerTransportProvider = HttpServletStreamableServerTransportProvider.builder() .objectMapper(new ObjectMapper()) + .contextExtractor(TEST_CONTEXT_EXTRACTOR) .mcpEndpoint(MESSAGE_ENDPOINT) .keepAliveInterval(Duration.ofSeconds(1)) .build(); @@ -54,7 +59,7 @@ public void before() { .put("httpclient", McpClient.sync(HttpClientStreamableHttpTransport.builder("http://localhost:" + PORT) .endpoint(MESSAGE_ENDPOINT) - .build()).initializationTimeout(Duration.ofHours(10)).requestTimeout(Duration.ofHours(10))); + .build()).requestTimeout(Duration.ofHours(10))); } @Override @@ -87,4 +92,9 @@ public void after() { protected void prepareClients(int port, String mcpEndpoint) { } + static McpTransportContextExtractor TEST_CONTEXT_EXTRACTOR = (r, tc) -> { + tc.put("important", "value"); + return tc; + }; + } diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/McpCompletionTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/McpCompletionTests.java index e6e80efb0..f915895be 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/McpCompletionTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/McpCompletionTests.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.server; import java.util.List; @@ -23,10 +27,12 @@ import io.modelcontextprotocol.spec.McpSchema; import io.modelcontextprotocol.spec.McpSchema.CompleteRequest; import io.modelcontextprotocol.spec.McpSchema.CompleteResult; +import io.modelcontextprotocol.spec.McpSchema.ErrorCodes; import io.modelcontextprotocol.spec.McpSchema.InitializeResult; import io.modelcontextprotocol.spec.McpSchema.Prompt; import io.modelcontextprotocol.spec.McpSchema.PromptArgument; import io.modelcontextprotocol.spec.McpSchema.ReadResourceResult; +import io.modelcontextprotocol.spec.McpSchema.Resource; import io.modelcontextprotocol.spec.McpSchema.ResourceReference; import io.modelcontextprotocol.spec.McpSchema.PromptReference; import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; @@ -80,7 +86,7 @@ public void after() { tomcat.destroy(); } catch (LifecycleException e) { - throw new RuntimeException("Failed to stop Tomcat", e); + e.printStackTrace(); } } } @@ -95,8 +101,13 @@ void testCompletionHandlerReceivesContext() { ResourceReference resourceRef = new ResourceReference("ref/resource", "test://resource/{param}"); - McpSchema.Resource resource = new McpSchema.Resource("test://resource/{param}", "Test Resource", - "A resource for testing", "text/plain", 123L, null); + var resource = Resource.builder() + .uri("test://resource/{param}") + .name("Test Resource") + .description("A resource for testing") + .mimeType("text/plain") + .size(123L) + .build(); var mcpServer = McpServer.sync(mcpServerTransportProvider) .capabilities(ServerCapabilities.builder().completions().build()) @@ -195,8 +206,13 @@ else if ("products_db".equals(db)) { return new CompleteResult(new CompleteResult.CompleteCompletion(List.of(), 0, false)); }; - McpSchema.Resource resource = new McpSchema.Resource("db://{database}/{table}", "Database Table", - "Resource representing a table in a database", "application/json", 456L, null); + McpSchema.Resource resource = Resource.builder() + .uri("db://{database}/{table}") + .name("Database Table") + .description("Resource representing a table in a database") + .mimeType("application/json") + .size(456L) + .build(); var mcpServer = McpServer.sync(mcpServerTransportProvider) .capabilities(ServerCapabilities.builder().completions().build()) @@ -250,7 +266,10 @@ void testCompletionErrorOnMissingContext() { // Check if database context is provided if (request.context() == null || request.context().arguments() == null || !request.context().arguments().containsKey("database")) { - throw new McpError("Please select a database first to see available tables"); + + throw McpError.builder(ErrorCodes.INVALID_REQUEST) + .message("Please select a database first to see available tables") + .build(); } // Normal completion if context is provided String db = request.context().arguments().get("database"); @@ -264,8 +283,13 @@ void testCompletionErrorOnMissingContext() { return new CompleteResult(new CompleteResult.CompleteCompletion(List.of(), 0, false)); }; - McpSchema.Resource resource = new McpSchema.Resource("db://{database}/{table}", "Database Table", - "Resource representing a table in a database", "application/json", 456L, null); + McpSchema.Resource resource = Resource.builder() + .uri("db://{database}/{table}") + .name("Database Table") + .description("Resource representing a table in a database") + .mimeType("application/json") + .size(456L) + .build(); var mcpServer = McpServer.sync(mcpServerTransportProvider) .capabilities(ServerCapabilities.builder().completions().build()) diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/McpServerProtocolVersionTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/McpServerProtocolVersionTests.java index 95086ee81..cdd2bacb7 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/McpServerProtocolVersionTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/McpServerProtocolVersionTests.java @@ -45,7 +45,9 @@ void shouldUseLatestVersionByDefault() { assertThat(jsonResponse.id()).isEqualTo(requestId); assertThat(jsonResponse.result()).isInstanceOf(McpSchema.InitializeResult.class); McpSchema.InitializeResult result = (McpSchema.InitializeResult) jsonResponse.result(); - assertThat(result.protocolVersion()).isEqualTo(transportProvider.protocolVersion()); + + var protocolVersion = transportProvider.protocolVersions().get(transportProvider.protocolVersions().size() - 1); + assertThat(result.protocolVersion()).isEqualTo(protocolVersion); server.closeGracefully().subscribe(); } @@ -93,7 +95,8 @@ void shouldSuggestLatestVersionForUnsupportedVersion() { assertThat(jsonResponse.id()).isEqualTo(requestId); assertThat(jsonResponse.result()).isInstanceOf(McpSchema.InitializeResult.class); McpSchema.InitializeResult result = (McpSchema.InitializeResult) jsonResponse.result(); - assertThat(result.protocolVersion()).isEqualTo(transportProvider.protocolVersion()); + var protocolVersion = transportProvider.protocolVersions().get(transportProvider.protocolVersions().size() - 1); + assertThat(result.protocolVersion()).isEqualTo(protocolVersion); server.closeGracefully().subscribe(); } diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/McpSyncServerExchangeTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/McpSyncServerExchangeTests.java index 63d827013..a73ec7209 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/McpSyncServerExchangeTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/McpSyncServerExchangeTests.java @@ -24,7 +24,6 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.never; -import static org.mockito.Mockito.reset; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/transport/HttpServletSseServerCustomContextPathTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/transport/HttpServletSseServerCustomContextPathTests.java index 2cd62889a..0462cbafe 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/transport/HttpServletSseServerCustomContextPathTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/transport/HttpServletSseServerCustomContextPathTests.java @@ -1,6 +1,7 @@ /* * Copyright 2024 - 2024 the original author or authors. */ + package io.modelcontextprotocol.server.transport; import com.fasterxml.jackson.databind.ObjectMapper; diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProviderIntegrationTests.java b/mcp/src/test/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProviderIntegrationTests.java deleted file mode 100644 index b04ecb3c4..000000000 --- a/mcp/src/test/java/io/modelcontextprotocol/server/transport/HttpServletSseServerTransportProviderIntegrationTests.java +++ /dev/null @@ -1,1390 +0,0 @@ -/* - * Copyright 2024 - 2025 the original author or authors. - */ -package io.modelcontextprotocol.server.transport; - -import java.time.Duration; -import java.util.ArrayList; -import java.util.List; -import java.util.Map; -import java.util.concurrent.CopyOnWriteArrayList; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Function; -import java.util.stream.Collectors; - -import com.fasterxml.jackson.databind.ObjectMapper; - -import io.modelcontextprotocol.client.McpClient; -import io.modelcontextprotocol.client.transport.HttpClientSseClientTransport; -import io.modelcontextprotocol.server.McpServer; -import io.modelcontextprotocol.server.McpServerFeatures; -import io.modelcontextprotocol.spec.McpError; -import io.modelcontextprotocol.spec.McpSchema; -import io.modelcontextprotocol.spec.McpSchema.CallToolResult; -import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities; -import io.modelcontextprotocol.spec.McpSchema.CreateMessageRequest; -import io.modelcontextprotocol.spec.McpSchema.CreateMessageResult; -import io.modelcontextprotocol.spec.McpSchema.ElicitRequest; -import io.modelcontextprotocol.spec.McpSchema.ElicitResult; -import io.modelcontextprotocol.spec.McpSchema.InitializeResult; -import io.modelcontextprotocol.spec.McpSchema.ModelPreferences; -import io.modelcontextprotocol.spec.McpSchema.Role; -import io.modelcontextprotocol.spec.McpSchema.Root; -import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities; -import io.modelcontextprotocol.spec.McpSchema.Tool; -import net.javacrumbs.jsonunit.core.Option; - -import org.apache.catalina.LifecycleException; -import org.apache.catalina.LifecycleState; -import org.apache.catalina.startup.Tomcat; -import org.junit.jupiter.api.AfterEach; -import org.junit.jupiter.api.BeforeEach; -import org.junit.jupiter.api.Test; -import reactor.core.publisher.Mono; -import reactor.test.StepVerifier; - -import org.springframework.web.client.RestClient; - -import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assertions.assertThatExceptionOfType; -import static org.assertj.core.api.InstanceOfAssertFactories.type; -import static org.awaitility.Awaitility.await; -import static org.mockito.Mockito.mock; -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; -import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json; - -class HttpServletSseServerTransportProviderIntegrationTests { - - private static final int PORT = TomcatTestUtil.findAvailablePort(); - - private static final String CUSTOM_SSE_ENDPOINT = "/somePath/sse"; - - private static final String CUSTOM_MESSAGE_ENDPOINT = "/otherPath/mcp/message"; - - private HttpServletSseServerTransportProvider mcpServerTransportProvider; - - McpClient.SyncSpec clientBuilder; - - private Tomcat tomcat; - - @BeforeEach - public void before() { - // Create and configure the transport provider - mcpServerTransportProvider = HttpServletSseServerTransportProvider.builder() - .objectMapper(new ObjectMapper()) - .messageEndpoint(CUSTOM_MESSAGE_ENDPOINT) - .sseEndpoint(CUSTOM_SSE_ENDPOINT) - .build(); - - tomcat = TomcatTestUtil.createTomcatServer("", PORT, mcpServerTransportProvider); - try { - tomcat.start(); - assertThat(tomcat.getServer().getState()).isEqualTo(LifecycleState.STARTED); - } - catch (Exception e) { - throw new RuntimeException("Failed to start Tomcat", e); - } - - this.clientBuilder = McpClient.sync(HttpClientSseClientTransport.builder("http://localhost:" + PORT) - .sseEndpoint(CUSTOM_SSE_ENDPOINT) - .build()); - } - - @AfterEach - public void after() { - if (mcpServerTransportProvider != null) { - mcpServerTransportProvider.closeGracefully().block(); - } - if (tomcat != null) { - try { - tomcat.stop(); - tomcat.destroy(); - } - catch (LifecycleException e) { - throw new RuntimeException("Failed to stop Tomcat", e); - } - } - } - - // --------------------------------------- - // Sampling Tests - // --------------------------------------- - @Test - // @Disabled - void testCreateMessageWithoutSamplingCapabilities() { - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - exchange.createMessage(mock(McpSchema.CreateMessageRequest.class)).block(); - - return Mono.just(mock(CallToolResult.class)); - }) - .build(); - - var server = McpServer.async(mcpServerTransportProvider).serverInfo("test-server", "1.0.0").tools(tool).build(); - - try ( - // Create client without sampling capabilities - var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample " + "client", "0.0.0")) - .build()) { - - assertThat(client.initialize()).isNotNull(); - - try { - client.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - } - catch (McpError e) { - assertThat(e).isInstanceOf(McpError.class) - .hasMessage("Client must be configured with sampling capabilities"); - } - } - server.close(); - } - - @Test - void testCreateMessageSuccess() { - - Function samplingHandler = request -> { - assertThat(request.messages()).hasSize(1); - assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); - - return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", - CreateMessageResult.StopReason.STOP_SEQUENCE); - }; - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var createMessageRequest = McpSchema.CreateMessageRequest.builder() - .messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, - new McpSchema.TextContent("Test message")))) - .modelPreferences(ModelPreferences.builder() - .hints(List.of()) - .costPriority(1.0) - .speedPriority(1.0) - .intelligencePriority(1.0) - .build()) - .build(); - - StepVerifier.create(exchange.createMessage(createMessageRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - } - mcpServer.close(); - } - - @Test - void testCreateMessageWithRequestTimeoutSuccess() throws InterruptedException { - - // Client - - Function samplingHandler = request -> { - assertThat(request.messages()).hasSize(1); - assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", - CreateMessageResult.StopReason.STOP_SEQUENCE); - }; - - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build(); - - // Server - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var craeteMessageRequest = McpSchema.CreateMessageRequest.builder() - .messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, - new McpSchema.TextContent("Test message")))) - .modelPreferences(ModelPreferences.builder() - .hints(List.of()) - .costPriority(1.0) - .speedPriority(1.0) - .intelligencePriority(1.0) - .build()) - .build(); - - StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(3)) - .tools(tool) - .build(); - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - - mcpClient.close(); - mcpServer.close(); - } - - @Test - void testCreateMessageWithRequestTimeoutFail() throws InterruptedException { - - // Client - - Function samplingHandler = request -> { - assertThat(request.messages()).hasSize(1); - assertThat(request.messages().get(0).content()).isInstanceOf(McpSchema.TextContent.class); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new CreateMessageResult(Role.USER, new McpSchema.TextContent("Test message"), "MockModelName", - CreateMessageResult.StopReason.STOP_SEQUENCE); - }; - - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().sampling().build()) - .sampling(samplingHandler) - .build(); - - // Server - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var craeteMessageRequest = McpSchema.CreateMessageRequest.builder() - .messages(List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, - new McpSchema.TextContent("Test message")))) - .modelPreferences(ModelPreferences.builder() - .hints(List.of()) - .costPriority(1.0) - .speedPriority(1.0) - .intelligencePriority(1.0) - .build()) - .build(); - - StepVerifier.create(exchange.createMessage(craeteMessageRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.role()).isEqualTo(Role.USER); - assertThat(result.content()).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content()).text()).isEqualTo("Test message"); - assertThat(result.model()).isEqualTo("MockModelName"); - assertThat(result.stopReason()).isEqualTo(CreateMessageResult.StopReason.STOP_SEQUENCE); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(1)) - .tools(tool) - .build(); - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThatExceptionOfType(McpError.class).isThrownBy(() -> { - mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - }).withMessageContaining("Timeout"); - - mcpClient.close(); - mcpServer.close(); - } - - // --------------------------------------- - // Elicitation Tests - // --------------------------------------- - @Test - // @Disabled - void testCreateElicitationWithoutElicitationCapabilities() { - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - exchange.createElicitation(mock(ElicitRequest.class)).block(); - - return Mono.just(mock(CallToolResult.class)); - }) - .build(); - - var server = McpServer.async(mcpServerTransportProvider).serverInfo("test-server", "1.0.0").tools(tool).build(); - - try ( - // Create client without elicitation capabilities - var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) { - - assertThat(client.initialize()).isNotNull(); - - try { - client.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - } - catch (McpError e) { - assertThat(e).isInstanceOf(McpError.class) - .hasMessage("Client must be configured with elicitation capabilities"); - } - } - server.closeGracefully().block(); - } - - @Test - void testCreateElicitationSuccess() { - - Function elicitationHandler = request -> { - assertThat(request.message()).isNotEmpty(); - assertThat(request.requestedSchema()).isNotNull(); - - return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); - }; - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var elicitationRequest = ElicitRequest.builder() - .message("Test message") - .requestedSchema( - Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) - .build(); - - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().elicitation().build()) - .elicitation(elicitationHandler) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - } - mcpServer.closeGracefully().block(); - } - - @Test - void testCreateElicitationWithRequestTimeoutSuccess() { - - // Client - - Function elicitationHandler = request -> { - assertThat(request.message()).isNotEmpty(); - assertThat(request.requestedSchema()).isNotNull(); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); - }; - - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().elicitation().build()) - .elicitation(elicitationHandler) - .build(); - - // Server - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var elicitationRequest = ElicitRequest.builder() - .message("Test message") - .requestedSchema( - Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) - .build(); - - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(3)) - .tools(tool) - .build(); - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - - mcpClient.closeGracefully(); - mcpServer.closeGracefully().block(); - } - - @Test - void testCreateElicitationWithRequestTimeoutFail() { - - // Client - - Function elicitationHandler = request -> { - assertThat(request.message()).isNotEmpty(); - assertThat(request.requestedSchema()).isNotNull(); - try { - TimeUnit.SECONDS.sleep(2); - } - catch (InterruptedException e) { - throw new RuntimeException(e); - } - return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); - }; - - var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) - .capabilities(ClientCapabilities.builder().elicitation().build()) - .elicitation(elicitationHandler) - .build(); - - // Server - - CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), - null); - - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - var elicitationRequest = ElicitRequest.builder() - .message("Test message") - .requestedSchema( - Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) - .build(); - - StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { - assertThat(result).isNotNull(); - assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); - assertThat(result.content().get("message")).isEqualTo("Test message"); - }).verifyComplete(); - - return Mono.just(callResponse); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .requestTimeout(Duration.ofSeconds(1)) - .tools(tool) - .build(); - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThatExceptionOfType(McpError.class).isThrownBy(() -> { - mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - }).withMessageContaining("Timeout"); - - mcpClient.closeGracefully(); - mcpServer.closeGracefully().block(); - } - - // --------------------------------------- - // Roots Tests - // --------------------------------------- - @Test - void testRootsSuccess() { - List roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2")); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(roots) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(rootsRef.get()).isNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(roots); - }); - - // Remove a root - mcpClient.removeRoot(roots.get(0).uri()); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(roots.get(1))); - }); - - // Add a new root - var root3 = new Root("uri3://", "root3"); - mcpClient.addRoot(root3); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(roots.get(1), root3)); - }); - - mcpServer.close(); - } - } - - @Test - void testRootsWithoutCapability() { - - McpServerFeatures.SyncToolSpecification tool = McpServerFeatures.SyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - exchange.listRoots(); // try to list roots - - return mock(CallToolResult.class); - }) - .build(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider).rootsChangeHandler((exchange, rootsUpdate) -> { - }).tools(tool).build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()).build()) { - - assertThat(mcpClient.initialize()).isNotNull(); - - // Attempt to list roots should fail - try { - mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - } - catch (McpError e) { - assertThat(e).isInstanceOf(McpError.class).hasMessage("Roots not supported"); - } - } - - mcpServer.close(); - } - - @Test - void testRootsNotificationWithEmptyRootsList() { - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(List.of()) // Empty roots list - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).isEmpty(); - }); - } - - mcpServer.close(); - } - - @Test - void testRootsWithMultipleHandlers() { - List roots = List.of(new Root("uri1://", "root1")); - - AtomicReference> rootsRef1 = new AtomicReference<>(); - AtomicReference> rootsRef2 = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef1.set(rootsUpdate)) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef2.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(roots) - .build()) { - - assertThat(mcpClient.initialize()).isNotNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef1.get()).containsAll(roots); - assertThat(rootsRef2.get()).containsAll(roots); - }); - } - - mcpServer.close(); - } - - @Test - void testRootsServerCloseWithActiveSubscription() { - List roots = List.of(new Root("uri1://", "root1")); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .rootsChangeHandler((exchange, rootsUpdate) -> rootsRef.set(rootsUpdate)) - .build(); - - try (var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) - .roots(roots) - .build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - mcpClient.rootsListChangedNotification(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(roots); - }); - } - - mcpServer.close(); - } - - // --------------------------------------- - // Tools Tests - // --------------------------------------- - - String emptyJsonSchema = """ - { - "$schema": "http://json-schema.org/draft-07/schema#", - "type": "object", - "properties": {} - } - """; - - @Test - void testToolCallSuccess() { - - var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - - McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - assertThat(McpTestServletFilter.getThreadLocalValue()).as("blocking code exectuion should be offloaded") - .isNull(); - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - return callResponse; - }) - .build(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool1) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response).isEqualTo(callResponse); - } - - mcpServer.close(); - } - - @Test - void testToolCallImmediateExecution() { - McpServerFeatures.SyncToolSpecification tool1 = new McpServerFeatures.SyncToolSpecification( - new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> { - var threadLocalValue = McpTestServletFilter.getThreadLocalValue(); - return CallToolResult.builder() - .addTextContent(threadLocalValue != null ? threadLocalValue : "") - .build(); - }); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool1) - .immediateExecution(true) - .build(); - - try (var mcpClient = clientBuilder.build()) { - mcpClient.initialize(); - - CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); - - assertThat(response).isNotNull(); - assertThat(response.content()).first() - .asInstanceOf(type(McpSchema.TextContent.class)) - .extracting(McpSchema.TextContent::text) - .isEqualTo(McpTestServletFilter.THREAD_LOCAL_VALUE); - } - - mcpServer.close(); - } - - @Test - void testToolListChangeHandlingSuccess() { - - var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null); - McpServerFeatures.SyncToolSpecification tool1 = McpServerFeatures.SyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema)) - .callHandler((exchange, request) -> { - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - return callResponse; - }) - .build(); - - AtomicReference> rootsRef = new AtomicReference<>(); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool1) - .build(); - - try (var mcpClient = clientBuilder.toolsChangeConsumer(toolsUpdate -> { - // perform a blocking call to a remote service - String response = RestClient.create() - .get() - .uri("https://raw.githubusercontent.com/modelcontextprotocol/java-sdk/refs/heads/main/README.md") - .retrieve() - .body(String.class); - assertThat(response).isNotBlank(); - rootsRef.set(toolsUpdate); - }).build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - assertThat(rootsRef.get()).isNull(); - - assertThat(mcpClient.listTools().tools()).contains(tool1.tool()); - - mcpServer.notifyToolsListChanged(); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool1.tool())); - }); - - // Remove a tool - mcpServer.removeTool("tool1"); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).isEmpty(); - }); - - // Add a new tool - McpServerFeatures.SyncToolSpecification tool2 = McpServerFeatures.SyncToolSpecification.builder() - .tool(new McpSchema.Tool("tool2", "tool2 description", emptyJsonSchema)) - .callHandler((exchange, request) -> callResponse) - .build(); - - mcpServer.addTool(tool2); - - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(rootsRef.get()).containsAll(List.of(tool2.tool())); - }); - } - - mcpServer.close(); - } - - @Test - void testInitialize() { - var mcpServer = McpServer.sync(mcpServerTransportProvider).build(); - - try (var mcpClient = clientBuilder.build()) { - - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - } - - mcpServer.close(); - } - - // --------------------------------------- - // Logging Tests - // --------------------------------------- - @Test - void testLoggingNotification() { - // Create a list to store received logging notifications - List receivedNotifications = new CopyOnWriteArrayList<>(); - - // Create server with a tool that sends logging notifications - McpServerFeatures.AsyncToolSpecification tool = McpServerFeatures.AsyncToolSpecification.builder() - .tool(new McpSchema.Tool("logging-test", "Test logging notifications", emptyJsonSchema)) - .callHandler((exchange, request) -> { - - // Create and send notifications with different levels - - // This should be filtered out (DEBUG < NOTICE) - exchange - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.DEBUG) - .logger("test-logger") - .data("Debug message") - .build()) - .block(); - - // This should be sent (NOTICE >= NOTICE) - exchange - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.NOTICE) - .logger("test-logger") - .data("Notice message") - .build()) - .block(); - - // This should be sent (ERROR > NOTICE) - exchange - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.ERROR) - .logger("test-logger") - .data("Error message") - .build()) - .block(); - - // This should be filtered out (INFO < NOTICE) - exchange - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.INFO) - .logger("test-logger") - .data("Another info message") - .build()) - .block(); - - // This should be sent (ERROR >= NOTICE) - exchange - .loggingNotification(McpSchema.LoggingMessageNotification.builder() - .level(McpSchema.LoggingLevel.ERROR) - .logger("test-logger") - .data("Another error message") - .build()) - .block(); - - return Mono.just(new CallToolResult("Logging test completed", false)); - }) - .build(); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().logging().tools(true).build()) - .tools(tool) - .build(); - try ( - // Create client with logging notification handler - var mcpClient = clientBuilder.loggingConsumer(notification -> { - receivedNotifications.add(notification); - }).build()) { - - // Initialize client - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Set minimum logging level to NOTICE - mcpClient.setLoggingLevel(McpSchema.LoggingLevel.NOTICE); - - // Call the tool that sends logging notifications - CallToolResult result = mcpClient.callTool(new McpSchema.CallToolRequest("logging-test", Map.of())); - assertThat(result).isNotNull(); - assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Logging test completed"); - - // Wait for notifications to be processed - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - - System.out.println("Received notifications: " + receivedNotifications); - - // Should have received 3 notifications (1 NOTICE and 2 ERROR) - assertThat(receivedNotifications).hasSize(3); - - Map notificationMap = receivedNotifications.stream() - .collect(Collectors.toMap(n -> n.data(), n -> n)); - - // First notification should be NOTICE level - assertThat(notificationMap.get("Notice message").level()).isEqualTo(McpSchema.LoggingLevel.NOTICE); - assertThat(notificationMap.get("Notice message").logger()).isEqualTo("test-logger"); - assertThat(notificationMap.get("Notice message").data()).isEqualTo("Notice message"); - - // Second notification should be ERROR level - assertThat(notificationMap.get("Error message").level()).isEqualTo(McpSchema.LoggingLevel.ERROR); - assertThat(notificationMap.get("Error message").logger()).isEqualTo("test-logger"); - assertThat(notificationMap.get("Error message").data()).isEqualTo("Error message"); - - // Third notification should be ERROR level - assertThat(notificationMap.get("Another error message").level()) - .isEqualTo(McpSchema.LoggingLevel.ERROR); - assertThat(notificationMap.get("Another error message").logger()).isEqualTo("test-logger"); - assertThat(notificationMap.get("Another error message").data()).isEqualTo("Another error message"); - }); - } - mcpServer.close(); - } - - // --------------------------------------- - // Progress Tests - // --------------------------------------- - @Test - void testProgressNotification() { - // Create a list to store received progress notifications - List receivedNotifications = new CopyOnWriteArrayList<>(); - - // Create server with a tool that sends progress notifications - McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification( - McpSchema.Tool.builder() - .name("progress-test") - .description("Test progress notifications") - .inputSchema(emptyJsonSchema) - .build(), - null, (exchange, request) -> { - - var progressToken = request.progressToken(); - - exchange - .progressNotification( - new McpSchema.ProgressNotification(progressToken, 0.1, 1.0, "Test progress 1/10")) - .block(); - - exchange - .progressNotification( - new McpSchema.ProgressNotification(progressToken, 0.5, 1.0, "Test progress 5/10")) - .block(); - - exchange - .progressNotification( - new McpSchema.ProgressNotification(progressToken, 1.0, 1.0, "Test progress 10/10")) - .block(); - - return Mono.just(new CallToolResult("Progress test completed", false)); - }); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().logging().tools(true).build()) - .tools(tool) - .build(); - - // Create client with progress notification handler - try (var mcpClient = clientBuilder.progressConsumer(receivedNotifications::add).build()) { - - // Initialize client - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call the tool that sends progress notifications - CallToolResult result = mcpClient.callTool( - new McpSchema.CallToolRequest("progress-test", Map.of(), Map.of("progressToken", "test-token"))); - assertThat(result).isNotNull(); - assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Progress test completed"); - - // Wait for notifications to be processed - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - // Should have received 3 notifications - assertThat(receivedNotifications).hasSize(3); - - // Check the progress notifications - assertThat(receivedNotifications.stream().map(McpSchema.ProgressNotification::progressToken)) - .containsExactlyInAnyOrder("test-token", "test-token", "test-token"); - assertThat(receivedNotifications.stream().map(McpSchema.ProgressNotification::progress)) - .containsExactlyInAnyOrder(0.1, 0.5, 1.0); - }); - } - finally { - mcpServer.close(); - } - } - - // --------------------------------------- - // Ping Tests - // --------------------------------------- - @Test - void testPingSuccess() { - // Create server with a tool that uses ping functionality - AtomicReference executionOrder = new AtomicReference<>(""); - - McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification( - new McpSchema.Tool("ping-async-test", "Test ping async behavior", emptyJsonSchema), - (exchange, request) -> { - - executionOrder.set(executionOrder.get() + "1"); - - // Test async ping behavior - return exchange.ping().doOnNext(result -> { - - assertThat(result).isNotNull(); - // Ping should return an empty object or map - assertThat(result).isInstanceOf(Map.class); - - executionOrder.set(executionOrder.get() + "2"); - assertThat(result).isNotNull(); - }).then(Mono.fromCallable(() -> { - executionOrder.set(executionOrder.get() + "3"); - return new CallToolResult("Async ping test completed", false); - })); - }); - - var mcpServer = McpServer.async(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - - // Initialize client - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call the tool that tests ping async behavior - CallToolResult result = mcpClient.callTool(new McpSchema.CallToolRequest("ping-async-test", Map.of())); - assertThat(result).isNotNull(); - assertThat(result.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - assertThat(((McpSchema.TextContent) result.content().get(0)).text()).isEqualTo("Async ping test completed"); - - // Verify execution order - assertThat(executionOrder.get()).isEqualTo("123"); - } - - mcpServer.close(); - } - - // --------------------------------------- - // Tool Structured Output Schema Tests - // --------------------------------------- - @Test - void testStructuredOutputValidationSuccess() { - // Create a tool with output schema - Map outputSchema = Map.of( - "type", "object", "properties", Map.of("result", Map.of("type", "number"), "operation", - Map.of("type", "string"), "timestamp", Map.of("type", "string")), - "required", List.of("result", "operation")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification tool = new McpServerFeatures.SyncToolSpecification(calculatorTool, - (exchange, request) -> { - String expression = (String) request.getOrDefault("expression", "2 + 3"); - double result = evaluateExpression(expression); - return CallToolResult.builder() - .structuredContent( - Map.of("result", result, "operation", expression, "timestamp", "2024-01-01T10:00:00Z")) - .build(); - }); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Verify tool is listed with output schema - var toolsList = mcpClient.listTools(); - assertThat(toolsList.tools()).hasSize(1); - assertThat(toolsList.tools().get(0).name()).isEqualTo("calculator"); - // Note: outputSchema might be null in sync server, but validation still works - - // Call tool with valid structured output - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isFalse(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - assertThatJson(((McpSchema.TextContent) response.content().get(0)).text()).when(Option.IGNORING_ARRAY_ORDER) - .when(Option.IGNORING_EXTRA_ARRAY_ITEMS) - .isObject() - .isEqualTo(json(""" - {"result":5.0,"operation":"2 + 3","timestamp":"2024-01-01T10:00:00Z"}""")); - - // Verify structured content (may be null in sync server but validation still - // works) - if (response.structuredContent() != null) { - assertThat(response.structuredContent()).containsEntry("result", 5.0) - .containsEntry("operation", "2 + 3") - .containsEntry("timestamp", "2024-01-01T10:00:00Z"); - } - } - - mcpServer.close(); - } - - @Test - void testStructuredOutputValidationFailure() { - - // Create a tool with output schema - Map outputSchema = Map.of("type", "object", "properties", - Map.of("result", Map.of("type", "number"), "operation", Map.of("type", "string")), "required", - List.of("result", "operation")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification tool = new McpServerFeatures.SyncToolSpecification(calculatorTool, - (exchange, request) -> { - // Return invalid structured output. Result should be number, missing - // operation - return CallToolResult.builder() - .addTextContent("Invalid calculation") - .structuredContent(Map.of("result", "not-a-number", "extra", "field")) - .build(); - }); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call tool with invalid structured output - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isTrue(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - String errorMessage = ((McpSchema.TextContent) response.content().get(0)).text(); - assertThat(errorMessage).contains("Validation failed"); - } - - mcpServer.close(); - } - - @Test - void testStructuredOutputMissingStructuredContent() { - // Create a tool with output schema - Map outputSchema = Map.of("type", "object", "properties", - Map.of("result", Map.of("type", "number")), "required", List.of("result")); - - Tool calculatorTool = Tool.builder() - .name("calculator") - .description("Performs mathematical calculations") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification tool = new McpServerFeatures.SyncToolSpecification(calculatorTool, - (exchange, request) -> { - // Return result without structured content but tool has output schema - return CallToolResult.builder().addTextContent("Calculation completed").build(); - }); - - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .tools(tool) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Call tool that should return structured content but doesn't - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isTrue(); - assertThat(response.content()).hasSize(1); - assertThat(response.content().get(0)).isInstanceOf(McpSchema.TextContent.class); - - String errorMessage = ((McpSchema.TextContent) response.content().get(0)).text(); - assertThat(errorMessage).isEqualTo( - "Response missing structured content which is expected when calling tool with non-empty outputSchema"); - } - - mcpServer.close(); - } - - @Test - void testStructuredOutputRuntimeToolAddition() { - // Start server without tools - var mcpServer = McpServer.sync(mcpServerTransportProvider) - .serverInfo("test-server", "1.0.0") - .capabilities(ServerCapabilities.builder().tools(true).build()) - .build(); - - try (var mcpClient = clientBuilder.build()) { - InitializeResult initResult = mcpClient.initialize(); - assertThat(initResult).isNotNull(); - - // Initially no tools - assertThat(mcpClient.listTools().tools()).isEmpty(); - - // Add tool with output schema at runtime - Map outputSchema = Map.of("type", "object", "properties", - Map.of("message", Map.of("type", "string"), "count", Map.of("type", "integer")), "required", - List.of("message", "count")); - - Tool dynamicTool = Tool.builder() - .name("dynamic-tool") - .description("Dynamically added tool") - .outputSchema(outputSchema) - .build(); - - McpServerFeatures.SyncToolSpecification toolSpec = new McpServerFeatures.SyncToolSpecification(dynamicTool, - (exchange, request) -> { - int count = (Integer) request.getOrDefault("count", 1); - return CallToolResult.builder() - .addTextContent("Dynamic tool executed " + count + " times") - .structuredContent(Map.of("message", "Dynamic execution", "count", count)) - .build(); - }); - - // Add tool to server - mcpServer.addTool(toolSpec); - - // Wait for tool list change notification - await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { - assertThat(mcpClient.listTools().tools()).hasSize(1); - }); - - // Verify tool was added with output schema - var toolsList = mcpClient.listTools(); - assertThat(toolsList.tools()).hasSize(1); - assertThat(toolsList.tools().get(0).name()).isEqualTo("dynamic-tool"); - // Note: outputSchema might be null in sync server, but validation still works - - // Call dynamically added tool - CallToolResult response = mcpClient - .callTool(new McpSchema.CallToolRequest("dynamic-tool", Map.of("count", 3))); - - assertThat(response).isNotNull(); - assertThat(response.isError()).isFalse(); - assertThat(response.structuredContent()).containsEntry("message", "Dynamic execution") - .containsEntry("count", 3); - } - - mcpServer.close(); - } - - private double evaluateExpression(String expression) { - // Simple expression evaluator for testing - return switch (expression) { - case "2 + 3" -> 5.0; - case "10 * 2" -> 20.0; - case "7 + 8" -> 15.0; - case "5 + 3" -> 8.0; - default -> 0.0; - }; - } - -} diff --git a/mcp/src/test/java/io/modelcontextprotocol/server/transport/TomcatTestUtil.java b/mcp/src/test/java/io/modelcontextprotocol/server/transport/TomcatTestUtil.java index 5a3928e02..2cf95dc94 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/server/transport/TomcatTestUtil.java +++ b/mcp/src/test/java/io/modelcontextprotocol/server/transport/TomcatTestUtil.java @@ -1,6 +1,7 @@ /* * Copyright 2025 - 2025 the original author or authors. */ + package io.modelcontextprotocol.server.transport; import java.io.IOException; diff --git a/mcp/src/test/java/io/modelcontextprotocol/spec/ArgumentException.java b/mcp/src/test/java/io/modelcontextprotocol/spec/ArgumentException.java index ba4e851f9..a0bd568ef 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/spec/ArgumentException.java +++ b/mcp/src/test/java/io/modelcontextprotocol/spec/ArgumentException.java @@ -1,3 +1,7 @@ +/* + * Copyright 2024-2025 the original author or authors. + */ + package io.modelcontextprotocol.spec; public class ArgumentException { diff --git a/mcp/src/test/java/io/modelcontextprotocol/spec/DefaultJsonSchemaValidatorTests.java b/mcp/src/test/java/io/modelcontextprotocol/spec/DefaultJsonSchemaValidatorTests.java index 9da31b38b..30158543d 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/spec/DefaultJsonSchemaValidatorTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/spec/DefaultJsonSchemaValidatorTests.java @@ -1,6 +1,7 @@ /* * Copyright 2024-2024 the original author or authors. */ + package io.modelcontextprotocol.spec; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -26,7 +27,6 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; -import io.modelcontextprotocol.spec.DefaultJsonSchemaValidator; import io.modelcontextprotocol.spec.JsonSchemaValidator.ValidationResponse; /** diff --git a/mcp/src/test/java/io/modelcontextprotocol/spec/McpSchemaTests.java b/mcp/src/test/java/io/modelcontextprotocol/spec/McpSchemaTests.java index fbbb4307e..a5b2137fd 100644 --- a/mcp/src/test/java/io/modelcontextprotocol/spec/McpSchemaTests.java +++ b/mcp/src/test/java/io/modelcontextprotocol/spec/McpSchemaTests.java @@ -1,6 +1,7 @@ /* * Copyright 2025 - 2025 the original author or authors. */ + package io.modelcontextprotocol.spec; import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; @@ -319,8 +320,8 @@ void testInitializeRequest() throws Exception { McpSchema.Implementation clientInfo = new McpSchema.Implementation("test-client", "1.0.0"); Map meta = Map.of("metaKey", "metaValue"); - McpSchema.InitializeRequest request = new McpSchema.InitializeRequest("2024-11-05", capabilities, clientInfo, - meta); + McpSchema.InitializeRequest request = new McpSchema.InitializeRequest(ProtocolVersions.MCP_2024_11_05, + capabilities, clientInfo, meta); String value = mapper.writeValueAsString(request); assertThatJson(value).when(Option.IGNORING_ARRAY_ORDER) @@ -342,8 +343,8 @@ void testInitializeResult() throws Exception { McpSchema.Implementation serverInfo = new McpSchema.Implementation("test-server", "1.0.0"); - McpSchema.InitializeResult result = new McpSchema.InitializeResult("2024-11-05", capabilities, serverInfo, - "Server initialized successfully"); + McpSchema.InitializeResult result = new McpSchema.InitializeResult(ProtocolVersions.MCP_2024_11_05, + capabilities, serverInfo, "Server initialized successfully"); String value = mapper.writeValueAsString(result); assertThatJson(value).when(Option.IGNORING_ARRAY_ORDER) diff --git a/mcp/src/test/java/io/modelcontextprotocol/spec/PromptReferenceEqualsTest.java b/mcp/src/test/java/io/modelcontextprotocol/spec/PromptReferenceEqualsTest.java new file mode 100644 index 000000000..25e22f968 --- /dev/null +++ b/mcp/src/test/java/io/modelcontextprotocol/spec/PromptReferenceEqualsTest.java @@ -0,0 +1,87 @@ +/* +* Copyright 2025 - 2025 the original author or authors. +*/ + +package io.modelcontextprotocol.spec; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.*; + +/** + * Test class to verify the equals method implementation for PromptReference. + */ +class PromptReferenceEqualsTest { + + @Test + void testEqualsWithSameIdentifierAndType() { + McpSchema.PromptReference ref1 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Test Title"); + McpSchema.PromptReference ref2 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Different Title"); + + assertTrue(ref1.equals(ref2), "PromptReferences with same identifier and type should be equal"); + assertEquals(ref1.hashCode(), ref2.hashCode(), "Equal objects should have same hash code"); + } + + @Test + void testEqualsWithDifferentIdentifier() { + McpSchema.PromptReference ref1 = new McpSchema.PromptReference("ref/prompt", "test-prompt-1", "Test Title"); + McpSchema.PromptReference ref2 = new McpSchema.PromptReference("ref/prompt", "test-prompt-2", "Test Title"); + + assertFalse(ref1.equals(ref2), "PromptReferences with different identifiers should not be equal"); + } + + @Test + void testEqualsWithDifferentType() { + McpSchema.PromptReference ref1 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Test Title"); + McpSchema.PromptReference ref2 = new McpSchema.PromptReference("ref/other", "test-prompt", "Test Title"); + + assertFalse(ref1.equals(ref2), "PromptReferences with different types should not be equal"); + } + + @Test + void testEqualsWithNull() { + McpSchema.PromptReference ref1 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Test Title"); + + assertFalse(ref1.equals(null), "PromptReference should not be equal to null"); + } + + @Test + void testEqualsWithDifferentClass() { + McpSchema.PromptReference ref1 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Test Title"); + String other = "not a PromptReference"; + + assertFalse(ref1.equals(other), "PromptReference should not be equal to different class"); + } + + @Test + void testEqualsWithSameInstance() { + McpSchema.PromptReference ref1 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Test Title"); + + assertTrue(ref1.equals(ref1), "PromptReference should be equal to itself"); + } + + @Test + void testEqualsIgnoresTitle() { + McpSchema.PromptReference ref1 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Title 1"); + McpSchema.PromptReference ref2 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Title 2"); + McpSchema.PromptReference ref3 = new McpSchema.PromptReference("ref/prompt", "test-prompt", null); + + assertTrue(ref1.equals(ref2), "PromptReferences should be equal regardless of title"); + assertTrue(ref1.equals(ref3), "PromptReferences should be equal even when one has null title"); + assertTrue(ref2.equals(ref3), "PromptReferences should be equal even when one has null title"); + } + + @Test + void testHashCodeConsistency() { + McpSchema.PromptReference ref1 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Test Title"); + McpSchema.PromptReference ref2 = new McpSchema.PromptReference("ref/prompt", "test-prompt", "Different Title"); + + assertEquals(ref1.hashCode(), ref2.hashCode(), "Objects that are equal should have the same hash code"); + + // Call hashCode multiple times to ensure consistency + int hashCode1 = ref1.hashCode(); + int hashCode2 = ref1.hashCode(); + assertEquals(hashCode1, hashCode2, "Hash code should be consistent across multiple calls"); + } + +} diff --git a/pom.xml b/pom.xml index b7a66aeec..2a637b9b7 100644 --- a/pom.xml +++ b/pom.xml @@ -6,7 +6,7 @@ io.modelcontextprotocol.sdk mcp-parent - 0.11.0-SNAPSHOT + 0.11.4-SNAPSHOT pom https://github.com/modelcontextprotocol/java-sdk