modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: Split transport implementations into modules

Browse source 
Split the  transport implementations into dedicated modules to improve modularity
and reduce core dependencies.

Added shared test module for common test utilities.

Key changes:

- Create mcp-transport/mcp-webflux-sse-transport module
 - Move SseClientTransport → WebFluxSseClientTransport
 - Move SseServerTransport → WebFluxSseServerTransport
 - Move related tests and utilities

- Create mcp-transport/mcp-webmvc-sse-transport module
 - Add new WebMvcSseServerTransport implementation
 - Add corresponding tests and integration tests

- Add framework-independent SSE client implementations to core module
 - Add FlowSseClient using java.net.http
 - Add HttpClientSseClientTransport using java.net.http

- Keep StdioClientTransport/StdioServerTransport in core (no external deps)

- Create mcp-test module for shared test infrastructure
 - Move abstract test base classes
 - Move shared test utilities and mocks
 - Update test dependencies

The modular approach allows users to choose between:

- Framework-specific implementations (WebFlux/WebMVC)
- Pure Java implementations (java.net.http)
- Basic implementations (STDIO)

Add breaking changes readme: 0.5.0-BREAKING-CHANGES.md

Resolves #54
This commit is contained in:
Christian Tzolov 2025-01-08 16:21:21 +01:00
commit 611f3258d6
17 changed files with 5002 additions and 1588 deletions
Download patch file Download diff file Expand all files Collapse all files

202
mcp/src/main/java/org/springframework/ai/mcp/client/transport/FlowSseClient.java Normal file
View file

@ -0,0 +1,202 @@
/*
* Copyright 2024 - 2024 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.
*/
package org.springframework.ai.mcp.client.transport;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Flow;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.regex.Pattern;
/**
* @author Christian Tzolov
*/
public class FlowSseClient {
private final HttpClient httpClient;
private static final Pattern EVENT_DATA_PATTERN = Pattern.compile("^data:(.+)$", Pattern.MULTILINE);
private static final Pattern EVENT_ID_PATTERN = Pattern.compile("^id:(.+)$", Pattern.MULTILINE);
private static final Pattern EVENT_TYPE_PATTERN = Pattern.compile("^event:(.+)$", Pattern.MULTILINE);
public FlowSseClient(HttpClient httpClient) {
this.httpClient = httpClient;
}
public void subscribe(String url, SSEEventHandler eventHandler) {
HttpRequest request = HttpRequest.newBuilder()
.uri(URI.create(url))
.header("Accept", "text/event-stream")
.header("Cache-Control", "no-cache")
.GET()
.build();
StringBuilder eventBuilder = new StringBuilder();
AtomicReference<String> currentEventId = new AtomicReference<>();
AtomicReference<String> currentEventType = new AtomicReference<>("message");
Flow.Subscriber<String> lineSubscriber = new Flow.Subscriber<>() {
private Flow.Subscription subscription;
@Override
public void onSubscribe(Flow.Subscription subscription) {
this.subscription = subscription;
subscription.request(Long.MAX_VALUE);
}
@Override
public void onNext(String line) {
if (line.isEmpty()) {
// Empty line means end of event
if (eventBuilder.length() > 0) {
String eventData = eventBuilder.toString();
SSEEvent event = new SSEEvent(currentEventId.get(), currentEventType.get(), eventData.trim());
eventHandler.onEvent(event);
eventBuilder.setLength(0);
}
}
else {
if (line.startsWith("data:")) {
var matcher = EVENT_DATA_PATTERN.matcher(line);
if (matcher.find()) {
eventBuilder.append(matcher.group(1).trim()).append("\n");
}
}
else if (line.startsWith("id:")) {
var matcher = EVENT_ID_PATTERN.matcher(line);
if (matcher.find()) {
currentEventId.set(matcher.group(1).trim());
}
}
else if (line.startsWith("event:")) {
var matcher = EVENT_TYPE_PATTERN.matcher(line);
if (matcher.find()) {
currentEventType.set(matcher.group(1).trim());
}
}
}
subscription.request(1);
}
@Override
public void onError(Throwable throwable) {
eventHandler.onError(throwable);
}
@Override
public void onComplete() {
// Handle any remaining event data
if (eventBuilder.length() > 0) {
String eventData = eventBuilder.toString();
SSEEvent event = new SSEEvent(currentEventId.get(), currentEventType.get(), eventData.trim());
eventHandler.onEvent(event);
}
}
};
Function<Flow.Subscriber<String>, HttpResponse.BodySubscriber<Void>> subscriberFactory = subscriber -> HttpResponse.BodySubscribers
.fromLineSubscriber(subscriber);
CompletableFuture<HttpResponse<Void>> future = this.httpClient.sendAsync(request,
info -> subscriberFactory.apply(lineSubscriber));
future.thenAccept(response -> {
if (response.statusCode() != 200) {
throw new RuntimeException("Failed to connect to SSE stream: " + response.statusCode());
}
}).exceptionally(throwable -> {
eventHandler.onError(throwable);
return null;
});
}
// Event handler interface
public interface SSEEventHandler {
void onEvent(SSEEvent event);
void onError(Throwable error);
}
// SSE Event class
public static class SSEEvent {
private final String id;
private final String type;
private final String data;
public SSEEvent(String id, String type, String data) {
this.id = id;
this.type = type;
this.data = data;
}
public String getId() {
return id;
}
public String getType() {
return type;
}
public String getData() {
return data;
}
@Override
public String toString() {
return "SSEEvent{" + "id='" + id + '\'' + ", type='" + type + '\'' + ", data='" + data + '\'' + '}';
}
}
public static void main(String[] args) {
FlowSseClient sseClient = new FlowSseClient(
HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build());
var eventHandler = new SSEEventHandler() {
@Override
public void onEvent(SSEEvent event) {
System.out.println("Received event: " + event);
}
@Override
public void onError(Throwable error) {
System.out.println("Error: " + error);
}
};
sseClient.subscribe("http://localhost:8080/sse", eventHandler);
try {
Thread.sleep(100000);
}
catch (InterruptedException e) {
e.printStackTrace();
}
}
}

192
mcp/src/main/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransport.java Normal file
View file

@ -0,0 +1,192 @@
/*
* Copyright 2024 - 2024 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.
*/
package org.springframework.ai.mcp.client.transport;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.mcp.client.transport.FlowSseClient.SSEEvent;
import org.springframework.ai.mcp.spec.ClientMcpTransport;
import org.springframework.ai.mcp.spec.McpError;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage;
import org.springframework.ai.mcp.util.Assert;
import reactor.core.publisher.Mono;
import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
/**
* Server-Sent Events (SSE) implementation of the
* {@link org.springframework.ai.mcp.spec.McpTransport} that follows the MCP HTTP with SSE
* transport specification, using Java's HttpClient.
*
* @author Christian Tzolov
*/
public class HttpClientSseClientTransport implements ClientMcpTransport {
private final static Logger logger = LoggerFactory.getLogger(HttpClientSseClientTransport.class);
private final static String MESSAGE_EVENT_TYPE = "message";
private final static String ENDPOINT_EVENT_TYPE = "endpoint";
private final static String SSE_ENDPOINT = "/sse";
private final String baseUri;
private final FlowSseClient sseClient;
private final HttpClient httpClient;
protected ObjectMapper objectMapper;
private volatile boolean isClosing = false;
private CountDownLatch closeLatch = new CountDownLatch(1);
private final AtomicReference<String> messageEndpoint = new AtomicReference<>();
private final AtomicReference<CompletableFuture<Void>> connectionFuture = new AtomicReference<>();
public HttpClientSseClientTransport(String baseUri) {
this(HttpClient.newBuilder(), baseUri, new ObjectMapper());
}
public HttpClientSseClientTransport(HttpClient.Builder clientBuilder, String baseUri, ObjectMapper objectMapper) {
Assert.notNull(objectMapper, "ObjectMapper must not be null");
Assert.notNull(clientBuilder, "clientBuilder must not be null");
this.baseUri = baseUri;
this.objectMapper = objectMapper;
this.httpClient = clientBuilder.connectTimeout(Duration.ofSeconds(10)).build();
this.sseClient = new FlowSseClient(this.httpClient);
}
@Override
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
CompletableFuture<Void> future = new CompletableFuture<>();
connectionFuture.set(future);
sseClient.subscribe(this.baseUri + SSE_ENDPOINT, new FlowSseClient.SSEEventHandler() {
@Override
public void onEvent(SSEEvent event) {
if (isClosing) {
return;
}
try {
if (ENDPOINT_EVENT_TYPE.equals(event.getType())) {
String endpoint = event.getData();
messageEndpoint.set(endpoint);
closeLatch.countDown();
future.complete(null);
}
else if (MESSAGE_EVENT_TYPE.equals(event.getType())) {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, event.getData());
handler.apply(Mono.just(message)).subscribe();
}
else {
logger.error("Received unrecognized SSE event type: " + event.getType());
}
}
catch (IOException e) {
logger.error("Error processing SSE event", e);
future.completeExceptionally(e);
}
}
@Override
public void onError(Throwable error) {
if (!isClosing) {
logger.error("SSE connection error", error);
future.completeExceptionally(error);
}
}
});
return Mono.fromFuture(future);
}
@Override
public Mono<Void> sendMessage(JSONRPCMessage message) {
if (isClosing) {
return Mono.empty();
}
try {
closeLatch.await(10, TimeUnit.SECONDS);
}
catch (InterruptedException e) {
return Mono.error(new McpError("Failed to wait for the message endpoint"));
}
String endpoint = messageEndpoint.get();
if (endpoint == null) {
return Mono.error(new McpError("No message endpoint available"));
}
try {
String jsonText = this.objectMapper.writeValueAsString(message);
HttpRequest request = HttpRequest.newBuilder()
.uri(URI.create(this.baseUri + endpoint))
.header("Content-Type", "application/json")
.POST(HttpRequest.BodyPublishers.ofString(jsonText))
.build();
return Mono.fromFuture(
httpClient.sendAsync(request, HttpResponse.BodyHandlers.discarding()).thenAccept(response -> {
if (response.statusCode() != 200) {
logger.error("Error sending message: {}", response.statusCode());
}
}));
}
catch (IOException e) {
if (!isClosing) {
return Mono.error(new RuntimeException("Failed to serialize message", e));
}
return Mono.empty();
}
}
@Override
public Mono<Void> closeGracefully() {
return Mono.fromRunnable(() -> {
isClosing = true;
CompletableFuture<Void> future = connectionFuture.get();
if (future != null && !future.isDone()) {
future.cancel(true);
}
});
}
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return this.objectMapper.convertValue(data, typeRef);
}
}

183
mcp/src/main/java/org/springframework/ai/mcp/client/transport/SSEClient.java Normal file
View file

@ -0,0 +1,183 @@
/*
* Copyright 2024 - 2024 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.
*/
package org.springframework.ai.mcp.client.transport;
/**
* @author Christian Tzolov
* @since 1.0.0
*/
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Pattern;
import java.util.stream.Stream;
public class SSEClient {
private final HttpClient httpClient;
private static final Pattern EVENT_DATA_PATTERN = Pattern.compile("^data:(.+)$", Pattern.MULTILINE);
private static final Pattern EVENT_ID_PATTERN = Pattern.compile("^id:(.+)$", Pattern.MULTILINE);
private static final Pattern EVENT_TYPE_PATTERN = Pattern.compile("^event:(.+)$", Pattern.MULTILINE);
public SSEClient() {
this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(10)).build();
}
public void subscribe(String url, SSEEventHandler eventHandler) {
HttpRequest request = HttpRequest.newBuilder()
.uri(URI.create(url))
.header("Accept", "text/event-stream")
.header("Cache-Control", "no-cache")
.GET()
.build();
CompletableFuture<Void> future = httpClient.sendAsync(request, HttpResponse.BodyHandlers.ofLines())
.thenAccept(response -> {
if (response.statusCode() == 200) {
processSSEStream(response, eventHandler);
}
else {
throw new RuntimeException("Failed to connect to SSE stream: " + response.statusCode());
}
});
// Handle errors
future.exceptionally(throwable -> {
eventHandler.onError(throwable);
return null;
});
}
private void processSSEStream(HttpResponse<Stream<String>> response, SSEEventHandler eventHandler) {
StringBuilder eventBuilder = new StringBuilder();
AtomicReference<String> currentEventId = new AtomicReference<>();
AtomicReference<String> currentEventType = new AtomicReference<>("message"); // default
// event
// type
response.body().forEach(line -> {
if (line.isEmpty()) {
// Empty line means end of event
if (eventBuilder.length() > 0) {
String eventData = eventBuilder.toString();
SSEEvent event = parseEvent(eventData, currentEventId.get(), currentEventType.get());
eventHandler.onEvent(event);
eventBuilder.setLength(0);
}
}
else {
if (line.startsWith("data:")) {
var matcher = EVENT_DATA_PATTERN.matcher(line);
if (matcher.find()) {
eventBuilder.append(matcher.group(1).trim()).append("\n");
}
}
else if (line.startsWith("id:")) {
var matcher = EVENT_ID_PATTERN.matcher(line);
if (matcher.find()) {
currentEventId.set(matcher.group(1).trim());
}
}
else if (line.startsWith("event:")) {
var matcher = EVENT_TYPE_PATTERN.matcher(line);
if (matcher.find()) {
currentEventType.set(matcher.group(1).trim());
}
}
}
});
}
private SSEEvent parseEvent(String eventData, String eventId, String eventType) {
return new SSEEvent(eventId, eventType, eventData.trim());
}
// Event handler interface
public interface SSEEventHandler {
void onEvent(SSEEvent event);
void onError(Throwable error);
}
// SSE Event class
public static class SSEEvent {
private final String id;
private final String type;
private final String data;
public SSEEvent(String id, String type, String data) {
this.id = id;
this.type = type;
this.data = data;
}
public String getId() {
return id;
}
public String getType() {
return type;
}
public String getData() {
return data;
}
@Override
public String toString() {
return "SSEEvent{" + "id='" + id + '\'' + ", type='" + type + '\'' + ", data='" + data + '\'' + '}';
}
}
public static void main(String[] args) {
SSEClient sseClient = new SSEClient();
var eventHandler = new SSEEventHandler() {
@Override
public void onEvent(SSEEvent event) {
System.out.println("Received event: " + event);
}
@Override
public void onError(Throwable error) {
System.out.println("Error: " + error);
}
};
sseClient.subscribe("http://localhost:8080/sse", eventHandler);
try {
Thread.sleep(100000);
}
catch (InterruptedException e) {
e.printStackTrace();
}
}
}

289
mcp/src/main/java/org/springframework/ai/mcp/client/transport/SseClientTransport.java
View file

@ -1,289 +0,0 @@
/*
* Copyright 2024 - 2024 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.
*/
package org.springframework.ai.mcp.client.transport;
import java.io.IOException;
import java.util.function.BiConsumer;
import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.core.publisher.SynchronousSink;
import reactor.core.scheduler.Schedulers;
import reactor.util.retry.Retry;
import reactor.util.retry.Retry.RetrySignal;
import org.springframework.ai.mcp.spec.ClientMcpTransport;
import org.springframework.ai.mcp.spec.McpError;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage;
import org.springframework.ai.mcp.util.Assert;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.reactive.function.client.WebClient;
/**
* Server-Sent Events (SSE) implementation of the
* {@link org.springframework.ai.mcp.spec.McpTransport} that follows the MCP HTTP with SSE
* transport specification.
*
* <p>
* This transport establishes a bidirectional communication channel where:
* <ul>
* <li>Inbound messages are received through an SSE connection from the server</li>
* <li>Outbound messages are sent via HTTP POST requests to a server-provided
* endpoint</li>
* </ul>
*
* <p>
* The message flow follows these steps:
* <ol>
* <li>The client establishes an SSE connection to the server's /sse endpoint</li>
* <li>The server sends an 'endpoint' event containing the URI for sending messages</li>
* </ol>
*
* This implementation handles automatic reconnection for transient failures and provides
* graceful shutdown capabilities. It uses {@link WebClient} for HTTP communications and
* supports JSON serialization/deserialization of messages.
*
* @author Christian Tzolov
* @see <a href=
* "https://spec.modelcontextprotocol.io/specification/basic/transports/#http-with-sse">MCP
* HTTP with SSE Transport Specification</a>
*/
public class SseClientTransport implements ClientMcpTransport {
private final static Logger logger = LoggerFactory.getLogger(SseClientTransport.class);
/**
* Event type for JSON-RPC messages received through the SSE connection. The server
* sends messages with this event type to transmit JSON-RPC protocol data.
*/
private final static String MESSAGE_EVENT_TYPE = "message";
/**
* Event type for receiving the message endpoint URI from the server. The server MUST
* send this event when a client connects, providing the URI where the client should
* send its messages via HTTP POST.
*/
private final static String ENDPOINT_EVENT_TYPE = "endpoint";
/**
* Default SSE endpoint path as specified by the MCP transport specification. This
* endpoint is used to establish the SSE connection with the server.
*/
private final static String SSE_ENDPOINT = "/sse";
/**
* Type reference for parsing SSE events containing string data.
*/
private final static ParameterizedTypeReference<ServerSentEvent<String>> SSE_TYPE = new ParameterizedTypeReference<>() {
};
/**
* WebClient instance for handling both SSE connections and HTTP POST requests. Used
* for establishing the SSE connection and sending outbound messages.
*/
private final WebClient webClient;
/**
* ObjectMapper for serializing outbound messages and deserializing inbound messages.
* Handles conversion between JSON-RPC messages and their string representation.
*/
protected ObjectMapper objectMapper;
/**
* Subscription for the SSE connection handling inbound messages. Used for cleanup
* during transport shutdown.
*/
private Disposable inboundSubscription;
/**
* Flag indicating if the transport is in the process of shutting down. Used to
* prevent new operations during shutdown and handle cleanup gracefully.
*/
private volatile boolean isClosing = false;
/**
* Sink for managing the message endpoint URI provided by the server. Stores the most
* recent endpoint URI and makes it available for outbound message processing.
*/
protected final Sinks.One<String> messageEndpointSink = Sinks.one();
/**
* Constructs a new SseClientTransport with the specified WebClient builder. Uses a
* default ObjectMapper instance for JSON processing.
* @param webClientBuilder the WebClient.Builder to use for creating the WebClient
* instance
* @throws IllegalArgumentException if webClientBuilder is null
*/
public SseClientTransport(WebClient.Builder webClientBuilder) {
this(webClientBuilder, new ObjectMapper());
}
/**
* Constructs a new SseClientTransport with the specified WebClient builder and
* ObjectMapper. Initializes both inbound and outbound message processing pipelines.
* @param webClientBuilder the WebClient.Builder to use for creating the WebClient
* instance
* @param objectMapper the ObjectMapper to use for JSON processing
* @throws IllegalArgumentException if either parameter is null
*/
public SseClientTransport(WebClient.Builder webClientBuilder, ObjectMapper objectMapper) {
Assert.notNull(objectMapper, "ObjectMapper must not be null");
Assert.notNull(webClientBuilder, "WebClient.Builder must not be null");
this.objectMapper = objectMapper;
this.webClient = webClientBuilder.build();
}
@Override
public Mono<Void> connect(Function<Mono<JSONRPCMessage>, Mono<JSONRPCMessage>> handler) {
Flux<ServerSentEvent<String>> events = eventStream();
this.inboundSubscription = events.concatMap(event -> Mono.just(event).<JSONRPCMessage>handle((e, s) -> {
if (ENDPOINT_EVENT_TYPE.equals(event.event())) {
String messageEndpointUri = event.data();
if (messageEndpointSink.tryEmitValue(messageEndpointUri).isSuccess()) {
s.complete();
}
else {
// TODO: clarify with the spec if multiple events can be
// received
s.error(new McpError("Failed to handle SSE endpoint event"));
}
}
else if (MESSAGE_EVENT_TYPE.equals(event.event())) {
try {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, event.data());
s.next(message);
}
catch (IOException ioException) {
s.error(ioException);
}
}
else {
s.error(new McpError("Received unrecognized SSE event type: " + event.event()));
}
}).transform(handler)).subscribe();
// The connection is established once the server sends the endpoint event
return messageEndpointSink.asMono().then();
}
@Override
public Mono<Void> sendMessage(JSONRPCMessage message) {
// The messageEndpoint is the endpoint URI to send the messages
// It is provided by the server as part of the endpoint event
return messageEndpointSink.asMono().flatMap(messageEndpointUri -> {
if (isClosing) {
return Mono.empty();
}
try {
String jsonText = this.objectMapper.writeValueAsString(message);
return webClient.post()
.uri(messageEndpointUri)
.contentType(MediaType.APPLICATION_JSON)
.bodyValue(jsonText)
.retrieve()
.toBodilessEntity()
.doOnSuccess(response -> {
logger.debug("Message sent successfully");
})
.doOnError(error -> {
if (!isClosing) {
logger.error("Error sending message: {}", error.getMessage());
}
});
}
catch (IOException e) {
if (!isClosing) {
return Mono.error(new RuntimeException("Failed to serialize message", e));
}
return Mono.empty();
}
}).then(); // TODO: Consider non-200-ok response
}
/**
* Initializes and starts the inbound SSE event processing. Establishes the SSE
* connection and sets up event handling for both message and endpoint events.
* Includes automatic retry logic for handling transient connection failures.
*/
// visible for tests
protected Flux<ServerSentEvent<String>> eventStream() {// @formatter:off
return this.webClient
.get()
.uri(SSE_ENDPOINT)
.accept(MediaType.TEXT_EVENT_STREAM)
.retrieve()
.bodyToFlux(SSE_TYPE)
.retryWhen(Retry.from(retrySignal -> retrySignal.handle(inboundRetryHandler)));
} // @formatter:on
/**
* Retry handler for the inbound SSE stream. Implements the retry logic for handling
* connection failures and other errors.
*/
private BiConsumer<RetrySignal, SynchronousSink<Object>> inboundRetryHandler = (retrySpec, sink) -> {
if (isClosing) {
logger.debug("SSE connection closed during shutdown");
sink.error(retrySpec.failure());
return;
}
if (retrySpec.failure() instanceof IOException) {
logger.debug("Retrying SSE connection after IO error");
sink.next(retrySpec);
return;
}
logger.error("Fatal SSE error, not retrying: {}", retrySpec.failure().getMessage());
sink.error(retrySpec.failure());
};
/**
* Implements graceful shutdown of the transport. Cleans up all resources including
* subscriptions and schedulers. Ensures orderly shutdown of both inbound and outbound
* message processing.
* @return a Mono that completes when shutdown is finished
*/
@Override
public Mono<Void> closeGracefully() { // @formatter:off
return Mono.fromRunnable(() -> {
isClosing = true;
// Dispose of subscriptions
if (inboundSubscription != null) {
inboundSubscription.dispose();
}
})
.then()
.subscribeOn(Schedulers.boundedElastic());
} // @formatter:on
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return this.objectMapper.convertValue(data, typeRef);
}
}

270
mcp/src/main/java/org/springframework/ai/mcp/server/transport/SseServerTransport.java
View file

@ -1,270 +0,0 @@
package org.springframework.ai.mcp.server.transport;
import java.io.IOException;
import java.time.Duration;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import org.springframework.ai.mcp.spec.McpError;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.ServerMcpTransport;
import org.springframework.ai.mcp.util.Assert;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.reactive.function.server.RouterFunction;
import org.springframework.web.reactive.function.server.RouterFunctions;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.reactive.function.server.ServerResponse;
/**
* Server-side implementation of the MCP HTTP with SSE transport specification. Provides
* endpoints for SSE connections and message handling.
*
* @author Christian Tzolov
*/
public class SseServerTransport implements ServerMcpTransport {
private final static Logger logger = LoggerFactory.getLogger(SseServerTransport.class);
/**
* Event type for JSON-RPC messages sent through the SSE connection.
*/
public final static String MESSAGE_EVENT_TYPE = "message";
/**
* Event type for sending the message endpoint URI to clients.
*/
public final static String ENDPOINT_EVENT_TYPE = "endpoint";
/**
* Default SSE endpoint path as specified by the MCP transport specification.
*/
public final static String SSE_ENDPOINT = "/sse";
private final ObjectMapper objectMapper;
private final String messageEndpoint;
private final RouterFunction<?> routerFunction;
/**
* Map of active client sessions, keyed by session ID.
*/
private final ConcurrentHashMap<String, ClientSession> sessions = new ConcurrentHashMap<>();
/**
* Flag indicating if the transport is shutting down.
*/
private volatile boolean isClosing = false;
private Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> connectHandler;
/**
* Constructs a new SseServerTransport.
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
* @param messageEndpoint The endpoint URI where clients should send messages
*/
public SseServerTransport(ObjectMapper objectMapper, String messageEndpoint) {
Assert.notNull(objectMapper, "ObjectMapper must not be null");
Assert.notNull(messageEndpoint, "Message endpoint must not be null");
this.objectMapper = objectMapper;
this.messageEndpoint = messageEndpoint;
this.routerFunction = RouterFunctions.route()
.GET(SSE_ENDPOINT, this::handleSseConnection)
.POST(messageEndpoint, this::handleMessage)
.build();
}
@Override
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
this.connectHandler = handler;
// Server-side transport doesn't initiate connections
return Mono.empty().then();
}
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
if (sessions.isEmpty()) {
logger.debug("No active sessions to broadcast message to");
return Mono.empty();
}
return Mono.<Void>create(sink -> {
try {// @formatter:off
String jsonText = objectMapper.writeValueAsString(message);
ServerSentEvent<Object> event = ServerSentEvent.builder()
.event(MESSAGE_EVENT_TYPE)
.data(jsonText)
.build();
logger.debug("Attempting to broadcast message to {} active sessions", sessions.size());
List<String> failedSessions = sessions.values().stream()
.filter(session -> session.messageSink.tryEmitNext(event).isFailure())
.map(session -> session.id)
.toList();
if (failedSessions.isEmpty()) {
logger.debug("Successfully broadcast message to all sessions");
sink.success();
}
else {
String error = "Failed to broadcast message to sessions: " + String.join(", ", failedSessions);
logger.error(error);
sink.error(new RuntimeException(error));
} // @formatter:on
}
catch (IOException e) {
logger.error("Failed to serialize message: {}", e.getMessage());
sink.error(e);
}
});
}
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return this.objectMapper.convertValue(data, typeRef);
}
@Override
public Mono<Void> closeGracefully() {
return Mono.fromRunnable(() -> {
isClosing = true;
logger.debug("Initiating graceful shutdown with {} active sessions", sessions.size());
}).then(Mono.when(sessions.values().stream().map(session -> {
String sessionId = session.id;
return Mono.fromRunnable(() -> session.close())
.then(Mono.delay(Duration.ofMillis(100)))
.then(Mono.fromRunnable(() -> sessions.remove(sessionId)));
}).toList()))
.timeout(Duration.ofSeconds(5))
.doOnSuccess(v -> logger.info("Graceful shutdown completed"))
.doOnError(e -> logger.error("Error during graceful shutdown: {}", e.getMessage()));
}
/**
* Get the router function for configuring the web server.
*/
public RouterFunction<?> getRouterFunction() {
return this.routerFunction;
}
/**
* Handles new SSE connection requests from clients.
*/
private Mono<ServerResponse> handleSseConnection(ServerRequest request) {
if (isClosing) {
return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).bodyValue("Server is shutting down");
}
String sessionId = UUID.randomUUID().toString();
logger.debug("Creating new SSE connection for session: {}", sessionId);
ClientSession session = new ClientSession(sessionId);
this.sessions.put(sessionId, session);
return ServerResponse.ok()
.contentType(MediaType.TEXT_EVENT_STREAM)
.body(Flux.<ServerSentEvent<?>>create(sink -> {
// Send initial endpoint event
logger.debug("Sending initial endpoint event to session: {}", sessionId);
sink.next(ServerSentEvent.builder().event(ENDPOINT_EVENT_TYPE).data(messageEndpoint).build());
// Subscribe to session messages
session.messageSink.asFlux()
.doOnSubscribe(s -> logger.debug("Session {} subscribed to message sink", sessionId))
.doOnComplete(() -> {
logger.debug("Session {} completed", sessionId);
sessions.remove(sessionId);
})
.doOnError(error -> {
logger.error("Error in session {}: {}", sessionId, error.getMessage());
sessions.remove(sessionId);
})
.doOnCancel(() -> {
logger.debug("Session {} cancelled", sessionId);
sessions.remove(sessionId);
})
.subscribe(event -> {
logger.debug("Forwarding event to session {}: {}", sessionId, event);
sink.next(event);
}, sink::error, sink::complete);
sink.onCancel(() -> {
logger.debug("Session {} cancelled", sessionId);
sessions.remove(sessionId);
});
}), ServerSentEvent.class);
}
/**
* Handles incoming messages from clients.
*/
private Mono<ServerResponse> handleMessage(ServerRequest request) {
if (isClosing) {
return ServerResponse.status(HttpStatus.SERVICE_UNAVAILABLE).bodyValue("Server is shutting down");
}
return request.bodyToMono(String.class).flatMap(body -> {
try {
McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body);
return Mono.just(message)
.transform(this.connectHandler)
.flatMap(response -> ServerResponse.ok().build())
.onErrorResume(error -> {
logger.error("Error processing message: {}", error.getMessage());
return ServerResponse.status(HttpStatus.INTERNAL_SERVER_ERROR)
.bodyValue(new McpError(error.getMessage()));
});
}
catch (IllegalArgumentException e) {
logger.error("Failed to deserialize message: {}", e.getMessage());
return ServerResponse.badRequest().bodyValue(new McpError("Invalid message format"));
}
catch (IOException e) {
logger.error("Failed to deserialize message: {}", e.getMessage());
return ServerResponse.badRequest().bodyValue(new McpError("Invalid message format"));
}
});
}
/**
* Represents an active client session.
*/
private static class ClientSession {
private final String id;
private final Sinks.Many<ServerSentEvent<?>> messageSink;
ClientSession(String id) {
this.id = id;
logger.debug("Creating new session: {}", id);
this.messageSink = Sinks.many().replay().latest();
logger.debug("Session {} initialized with replay sink", id);
}
void close() {
logger.debug("Closing session: {}", id);
Sinks.EmitResult result = messageSink.tryEmitComplete();
if (result.isFailure()) {
logger.warn("Failed to complete message sink for session {}: {}", id, result);
}
else {
logger.debug("Successfully completed message sink for session {}", id);
}
}
}
}

6
mcp/src/main/java/org/springframework/ai/mcp/util/Utils.java
View file

@ -21,11 +21,9 @@ import java.util.Map;
import reactor.util.annotation.Nullable;
import org.springframework.lang.Contract;
/**
* Miscellaneous utility methods.
*
*
* @author Christian Tzolov
*/
@ -52,7 +50,6 @@ public final class Utils {
* @param collection the Collection to check
* @return whether the given Collection is empty
*/
@Contract("null -> true")
public static boolean isEmpty(@Nullable Collection<?> collection) {
return (collection == null || collection.isEmpty());
}
@ -63,7 +60,6 @@ public final class Utils {
* @param map the Map to check
* @return whether the given Map is empty
*/
@Contract("null -> true")
public static boolean isEmpty(@Nullable Map<?, ?> map) {
return (map == null || map.isEmpty());
}

61
mcp/src/test/java/org/springframework/ai/mcp/attic/_DemoServer.java_ Normal file
View file

@ -0,0 +1,61 @@
/*
* Copyright 2024 - 2024 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.
*/
package org.springframework.ai.mcp.attic;
import java.util.List;
import java.util.Map;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.netty.http.server.HttpServer;
import org.springframework.ai.mcp.server.McpServer;
import org.springframework.ai.mcp.server.transport.SseServerTransport;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
import org.springframework.http.server.reactive.HttpHandler;
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
import org.springframework.web.reactive.function.server.RouterFunctions;
/**
*
* https://docs.spring.io/spring-framework/reference/web/webflux-functional.html#webflux-fn-running
* https://docs.spring.io/spring-framework/reference/web/webflux/reactive-spring.html#webflux-httphandler
*
* @author Christian Tzolov
*/
public class DemoServer {
public static void main(String[] args) {
SseServerTransport transport = new SseServerTransport(new ObjectMapper(), "/mcp/message");
var mcpServer = McpServer.using(transport)
.serverInfo("Weather Forecast", "1.0.0")
.tool(new McpSchema.Tool("weather", "Weather forecast tool by location", "{ \"type\": \"object\" }"),
(arguments) -> {
String city = (String) arguments.get("city");
return new CallToolResult(List.of(), false);
})
.async();
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
HttpServer httpServer = HttpServer.create().port(8080).handle(adapter);
httpServer.bindNow().onDispose().block();
}
}

297
mcp/src/test/java/org/springframework/ai/mcp/attic/_HttpServletSseServerTransport.java_ Normal file
View file

@ -0,0 +1,297 @@
/*
* Copyright 2024 - 2024 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.
*/
package org.springframework.ai.mcp.server.transport;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.servlet.AsyncContext;
import jakarta.servlet.ServletException;
import jakarta.servlet.annotation.WebServlet;
import jakarta.servlet.http.HttpServlet;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
import org.springframework.ai.mcp.spec.McpError;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.ServerMcpTransport;
/**
* A Servlet-based implementation of the MCP HTTP with SSE transport specification. This
* implementation provides similar functionality to FunctionalSseServerTransport but uses
* the traditional Servlet API instead of WebFlux.
*
* @author Christian Tzolov
* @since 1.0.0
*/
@WebServlet(asyncSupported = true)
public class HttpServletSseServerTransport extends HttpServlet implements ServerMcpTransport {
private static final Logger logger = LoggerFactory.getLogger(HttpServletSseServerTransport.class);
public static final String SSE_ENDPOINT = "/sse";
public static final String MESSAGE_EVENT_TYPE = "message";
public static final String ENDPOINT_EVENT_TYPE = "endpoint";
private final ObjectMapper objectMapper;
private final String messageEndpoint;
private final Map<String, ClientSession> sessions = new ConcurrentHashMap<>();
private final AtomicBoolean isClosing = new AtomicBoolean(false);
private Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> connectHandler;
public HttpServletSseServerTransport(ObjectMapper objectMapper, String messageEndpoint) {
this.objectMapper = objectMapper;
this.messageEndpoint = messageEndpoint;
}
@Override
protected void doGet(HttpServletRequest request, HttpServletResponse response)
throws ServletException, IOException {
String pathInfo = request.getPathInfo();
if (!SSE_ENDPOINT.equals(pathInfo)) {
response.sendError(HttpServletResponse.SC_NOT_FOUND);
return;
}
if (isClosing.get()) {
response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "Server is shutting down");
return;
}
response.setContentType("text/event-stream");
response.setCharacterEncoding("UTF-8");
response.setHeader("Cache-Control", "no-cache");
response.setHeader("Connection", "keep-alive");
response.setHeader("Access-Control-Allow-Origin", "*");
String sessionId = UUID.randomUUID().toString();
AsyncContext asyncContext = request.startAsync();
asyncContext.setTimeout(0);
PrintWriter writer = response.getWriter();
ClientSession session = new ClientSession(sessionId, asyncContext, writer);
sessions.put(sessionId, session);
// Send initial endpoint event
sendEvent(writer, ENDPOINT_EVENT_TYPE, messageEndpoint);
}
@Override
protected void doPost(HttpServletRequest request, HttpServletResponse response)
throws ServletException, IOException {
if (isClosing.get()) {
response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "Server is shutting down");
return;
}
String pathInfo = request.getPathInfo();
if (!messageEndpoint.equals(pathInfo)) {
response.sendError(HttpServletResponse.SC_NOT_FOUND);
return;
}
try {
BufferedReader reader = request.getReader();
StringBuilder body = new StringBuilder();
String line;
while ((line = reader.readLine()) != null) {
body.append(line);
}
McpSchema.JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(objectMapper, body.toString());
if (connectHandler != null) {
connectHandler.apply(Mono.just(message)).subscribe(responseMessage -> {
try {
response.setContentType("application/json");
response.setCharacterEncoding("UTF-8");
String jsonResponse = objectMapper.writeValueAsString(responseMessage);
PrintWriter writer = response.getWriter();
writer.write(jsonResponse);
writer.flush();
}
catch (Exception e) {
logger.error("Error sending response: {}", e.getMessage());
try {
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
"Error processing response: " + e.getMessage());
}
catch (IOException ex) {
logger.error("Failed to send error response: {}", ex.getMessage());
}
}
}, error -> {
try {
logger.error("Error processing message: {}", error.getMessage());
McpError mcpError = new McpError(error.getMessage());
response.setContentType("application/json");
response.setCharacterEncoding("UTF-8");
response.setStatus(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
String jsonError = objectMapper.writeValueAsString(mcpError);
PrintWriter writer = response.getWriter();
writer.write(jsonError);
writer.flush();
}
catch (IOException e) {
logger.error("Failed to send error response: {}", e.getMessage());
try {
response.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR,
"Error sending error response: " + e.getMessage());
}
catch (IOException ex) {
logger.error("Failed to send error response: {}", ex.getMessage());
}
}
});
}
else {
response.sendError(HttpServletResponse.SC_SERVICE_UNAVAILABLE, "No message handler configured");
}
}
catch (Exception e) {
logger.error("Invalid message format: {}", e.getMessage());
try {
McpError mcpError = new McpError("Invalid message format: " + e.getMessage());
response.setContentType("application/json");
response.setCharacterEncoding("UTF-8");
response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
String jsonError = objectMapper.writeValueAsString(mcpError);
PrintWriter writer = response.getWriter();
writer.write(jsonError);
writer.flush();
}
catch (IOException ex) {
logger.error("Failed to send error response: {}", ex.getMessage());
response.sendError(HttpServletResponse.SC_BAD_REQUEST, "Invalid message format");
}
}
}
@Override
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
this.connectHandler = handler;
return Mono.empty();
}
@Override
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
if (sessions.isEmpty()) {
logger.debug("No active sessions to broadcast message to");
return Mono.empty();
}
return Mono.create(sink -> {
try {
String jsonText = objectMapper.writeValueAsString(message);
sessions.values().forEach(session -> {
try {
sendEvent(session.writer, MESSAGE_EVENT_TYPE, jsonText);
}
catch (IOException e) {
logger.error("Failed to send message to session {}: {}", session.id, e.getMessage());
removeSession(session);
}
});
sink.success();
}
catch (Exception e) {
logger.error("Failed to process message: {}", e.getMessage());
sink.error(new McpError("Failed to process message: " + e.getMessage()));
}
});
}
@Override
public void close() {
ServerMcpTransport.super.close();
}
@Override
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
return objectMapper.convertValue(data, typeRef);
}
@Override
public Mono<Void> closeGracefully() {
isClosing.set(true);
logger.debug("Initiating graceful shutdown with {} active sessions", sessions.size());
return Mono.create(sink -> {
sessions.values().forEach(this::removeSession);
sink.success();
});
}
private void sendEvent(PrintWriter writer, String eventType, String data) throws IOException {
writer.write("event: " + eventType + "\n");
writer.write("data: " + data + "\n\n");
writer.flush();
if (writer.checkError()) {
throw new IOException("Client disconnected");
}
}
private void removeSession(ClientSession session) {
sessions.remove(session.id);
session.asyncContext.complete();
}
private static class ClientSession {
private final String id;
private final AsyncContext asyncContext;
private final PrintWriter writer;
ClientSession(String id, AsyncContext asyncContext, PrintWriter writer) {
this.id = id;
this.asyncContext = asyncContext;
this.writer = writer;
}
}
@Override
public void destroy() {
closeGracefully().block();
super.destroy();
}
}

81
mcp/src/test/java/org/springframework/ai/mcp/attic/_MainSSE.java_ Normal file
View file

@ -0,0 +1,81 @@
/*
* Copyright 2024 - 2024 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.
*/
package org.springframework.ai.mcp.attic;
import java.time.Duration;
import java.util.Map;
import org.springframework.ai.mcp.client.McpClient;
import org.springframework.ai.mcp.client.McpSyncClient;
import org.springframework.ai.mcp.client.transport.SseClientTransport;
import org.springframework.ai.mcp.spec.McpSchema.CallToolRequest;
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
import org.springframework.ai.mcp.spec.McpSchema.ListPromptsResult;
import org.springframework.ai.mcp.spec.McpSchema.ListResourcesResult;
import org.springframework.ai.mcp.spec.McpSchema.ListToolsResult;
import org.springframework.ai.mcp.spec.McpSchema.Resource;
import org.springframework.web.reactive.function.client.WebClient;
/**
* @author Christian Tzolov
* @since 1.0.0
*/
public class MainSSE {
public static void main(String[] args) {
try (McpSyncClient mcpClient = McpClient
.using(new SseClientTransport(WebClient.builder().baseUrl("http://localhost:3001")))
.requestTimeout(Duration.ofSeconds(1000))
.sync()) {
mcpClient.initialize();
ListToolsResult tools = mcpClient.listTools();
System.out.println("Tools: " + tools);
mcpClient.ping();
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", "Hello MCP Spring AI!"));
CallToolResult callToolResult = mcpClient.callTool(callToolRequest);
System.out.println("Call Tool Result: " + callToolResult);
// mcpClient.sendRootsListChanged();
// Resources
ListResourcesResult resources = mcpClient.listResources();
System.out.println("Resources Size: " + resources.resources().size());
System.out.println("Resources: " + resources);
for (Resource resource : resources.resources()) {
System.out.println(mcpClient.readResource(resource));
}
var resourceTemplate = mcpClient.listResourceTemplates();
System.out.println("Resource Templates: " + resourceTemplate);
ListPromptsResult prompts = mcpClient.listPrompts();
for (var prompt : prompts.prompts()) {
System.out.println("Prompt: " + prompt);
}
}
catch (Exception e) {
e.printStackTrace();
}
}
}

61
mcp/src/test/java/org/springframework/ai/mcp/client/SseMcpAsyncClientTests.java
View file

@ -1,61 +0,0 @@
/*
* Copyright 2024-2024 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.
*/
package org.springframework.ai.mcp.client;
import org.junit.jupiter.api.Timeout;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.springframework.ai.mcp.client.transport.SseClientTransport;
import org.springframework.ai.mcp.spec.ClientMcpTransport;
import org.springframework.web.reactive.function.client.WebClient;
/**
* Tests for the {@link McpAsyncClient} with {@link SseClientTransport}.
*
* @author Christian Tzolov
*/
@Timeout(15) // Giving extra time beyond the client timeout
class SseMcpAsyncClientTests extends AbstractMcpAsyncClientTests {
static String host = "http://localhost:3001";
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
@SuppressWarnings("resource")
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
.withExposedPorts(3001)
.waitingFor(Wait.forHttp("/").forStatusCode(404));
@Override
protected ClientMcpTransport createMcpTransport() {
return new SseClientTransport(WebClient.builder().baseUrl(host));
}
@Override
protected void onStart() {
container.start();
int port = container.getMappedPort(3001);
host = "http://" + container.getHost() + ":" + port;
}
@Override
public void onClose() {
container.stop();
}
}

61
mcp/src/test/java/org/springframework/ai/mcp/client/SseMcpSyncClientTests.java
View file

@ -1,61 +0,0 @@
/*
* Copyright 2024-2024 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.
*/
package org.springframework.ai.mcp.client;
import org.junit.jupiter.api.Timeout;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.springframework.ai.mcp.client.transport.SseClientTransport;
import org.springframework.ai.mcp.spec.ClientMcpTransport;
import org.springframework.web.reactive.function.client.WebClient;
/**
* Tests for the {@link McpSyncClient} with {@link SseClientTransport}.
*
* @author Christian Tzolov
*/
@Timeout(15) // Giving extra time beyond the client timeout
class SseMcpSyncClientTests extends AbstractMcpSyncClientTests {
static String host = "http://localhost:3001";
// Uses the https://github.com/tzolov/mcp-everything-server-docker-image
@SuppressWarnings("resource")
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
.withExposedPorts(3001)
.waitingFor(Wait.forHttp("/").forStatusCode(404));
@Override
protected ClientMcpTransport createMcpTransport() {
return new SseClientTransport(WebClient.builder().baseUrl(host));
}
@Override
protected void onStart() {
container.start();
int port = container.getMappedPort(3001);
host = "http://" + container.getHost() + ":" + port;
}
@Override
protected void onClose() {
container.stop();
}
}

290
mcp/src/test/java/org/springframework/ai/mcp/client/transport/HttpClientSseClientTransportTests.java Normal file
View file

@ -0,0 +1,290 @@
/*
* Copyright 2024-2024 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.
*/
package org.springframework.ai.mcp.client.transport;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.Timeout;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.containers.wait.strategy.Wait;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.test.StepVerifier;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
import org.springframework.http.codec.ServerSentEvent;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatCode;
/**
* Tests for the {@link HttpClientSseClientTransport} class.
*
* @author Christian Tzolov
*/
@Timeout(15)
class HttpClientSseClientTransportTests {
static String host = "http://localhost:3001";
@SuppressWarnings("resource")
GenericContainer<?> container = new GenericContainer<>("docker.io/tzolov/mcp-everything-server:v1")
.withLogConsumer(outputFrame -> System.out.println(outputFrame.getUtf8String()))
.withExposedPorts(3001)
.waitingFor(Wait.forHttp("/").forStatusCode(404));
private TestHttpClientSseClientTransport transport;
// Test class to access protected methods
static class TestHttpClientSseClientTransport extends HttpClientSseClientTransport {
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
private Sinks.Many<ServerSentEvent<String>> events = Sinks.many().unicast().onBackpressureBuffer();
public TestHttpClientSseClientTransport(String baseUri) {
super(baseUri);
}
public int getInboundMessageCount() {
return inboundMessageCount.get();
}
public void simulateEndpointEvent(String jsonMessage) {
events.tryEmitNext(ServerSentEvent.<String>builder().event("endpoint").data(jsonMessage).build());
inboundMessageCount.incrementAndGet();
}
public void simulateMessageEvent(String jsonMessage) {
events.tryEmitNext(ServerSentEvent.<String>builder().event("message").data(jsonMessage).build());
inboundMessageCount.incrementAndGet();
}
}
void startContainer() {
container.start();
int port = container.getMappedPort(3001);
host = "http://" + container.getHost() + ":" + port;
}
@BeforeEach
void setUp() {
startContainer();
transport = new TestHttpClientSseClientTransport(host);
transport.connect(Function.identity()).block();
}
@AfterEach
void afterEach() {
if (transport != null) {
assertThatCode(() -> transport.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();
}
cleanup();
}
void cleanup() {
container.stop();
}
@Test
void testMessageProcessing() {
// Create a test message
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
Map.of("key", "value"));
// Simulate receiving the message
transport.simulateMessageEvent("""
{
"jsonrpc": "2.0",
"method": "test-method",
"id": "test-id",
"params": {"key": "value"}
}
""");
// Subscribe to messages and verify
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
}
@Test
void testResponseMessageProcessing() {
// Simulate receiving a response message
transport.simulateMessageEvent("""
{
"jsonrpc": "2.0",
"id": "test-id",
"result": {"status": "success"}
}
""");
// Create and send a request message
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
Map.of("key", "value"));
// Verify message handling
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
}
@Test
void testErrorMessageProcessing() {
// Simulate receiving an error message
transport.simulateMessageEvent("""
{
"jsonrpc": "2.0",
"id": "test-id",
"error": {
"code": -32600,
"message": "Invalid Request"
}
}
""");
// Create and send a request message
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
Map.of("key", "value"));
// Verify message handling
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
}
@Test
void testNotificationMessageProcessing() {
// Simulate receiving a notification message (no id)
transport.simulateMessageEvent("""
{
"jsonrpc": "2.0",
"method": "update",
"params": {"status": "processing"}
}
""");
// Verify the notification was processed
assertThat(transport.getInboundMessageCount()).isEqualTo(1);
}
@Test
void testGracefulShutdown() {
// Test graceful shutdown
StepVerifier.create(transport.closeGracefully()).verifyComplete();
// Create a test message
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
Map.of("key", "value"));
// Verify message is not processed after shutdown
StepVerifier.create(transport.sendMessage(testMessage)).verifyComplete();
// Message count should remain 0 after shutdown
assertThat(transport.getInboundMessageCount()).isEqualTo(0);
}
@Test
void testRetryBehavior() {
// Create a client that simulates connection failures
HttpClientSseClientTransport failingTransport = new HttpClientSseClientTransport("http://non-existent-host");
// Verify that the transport attempts to reconnect
StepVerifier.create(Mono.delay(Duration.ofSeconds(2))).expectNextCount(1).verifyComplete();
// Clean up
failingTransport.closeGracefully().block();
}
@Test
void testMultipleMessageProcessing() {
// Simulate receiving multiple messages in sequence
transport.simulateMessageEvent("""
{
"jsonrpc": "2.0",
"method": "method1",
"id": "id1",
"params": {"key": "value1"}
}
""");
transport.simulateMessageEvent("""
{
"jsonrpc": "2.0",
"method": "method2",
"id": "id2",
"params": {"key": "value2"}
}
""");
// Create and send corresponding messages
JSONRPCRequest message1 = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "method1", "id1",
Map.of("key", "value1"));
JSONRPCRequest message2 = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "method2", "id2",
Map.of("key", "value2"));
// Verify both messages are processed
StepVerifier.create(transport.sendMessage(message1).then(transport.sendMessage(message2))).verifyComplete();
// Verify message count
assertThat(transport.getInboundMessageCount()).isEqualTo(2);
}
@Test
void testMessageOrderPreservation() {
// Simulate receiving messages in a specific order
transport.simulateMessageEvent("""
{
"jsonrpc": "2.0",
"method": "first",
"id": "1",
"params": {"sequence": 1}
}
""");
transport.simulateMessageEvent("""
{
"jsonrpc": "2.0",
"method": "second",
"id": "2",
"params": {"sequence": 2}
}
""");
transport.simulateMessageEvent("""
{
"jsonrpc": "2.0",
"method": "third",
"id": "3",
"params": {"sequence": 3}
}
""");
// Verify message count and order
assertThat(transport.getInboundMessageCount()).isEqualTo(3);
}
}

1
mcp/src/test/java/org/springframework/ai/mcp/server/AbstractMcpAsyncServerTests.java
View file

@ -18,7 +18,6 @@ package org.springframework.ai.mcp.server;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;

66
mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpAsyncServerTests.java
View file

@ -1,66 +0,0 @@
/*
* Copyright 2024-2024 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.
*/
package org.springframework.ai.mcp.server;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Timeout;
import reactor.netty.DisposableServer;
import reactor.netty.http.server.HttpServer;
import org.springframework.ai.mcp.server.transport.SseServerTransport;
import org.springframework.ai.mcp.spec.ServerMcpTransport;
import org.springframework.http.server.reactive.HttpHandler;
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
import org.springframework.web.reactive.function.server.RouterFunctions;
/**
* Tests for {@link McpAsyncServer} using {@link SseServerTransport}.
*
* @author Christian Tzolov
*/
@Timeout(15) // Giving extra time beyond the client timeout
class SseMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
private static final int PORT = 8181;
private static final String MESSAGE_ENDPOINT = "/mcp/message";
private DisposableServer httpServer;
@Override
protected ServerMcpTransport createMcpTransport() {
var transport = new SseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
return transport;
}
@Override
protected void onStart() {
}
@Override
protected void onClose() {
if (httpServer != null) {
httpServer.disposeNow();
}
}
}

66
mcp/src/test/java/org/springframework/ai/mcp/server/SseMcpSyncServerTests.java
View file

@ -1,66 +0,0 @@
/*
* Copyright 2024-2024 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.
*/
package org.springframework.ai.mcp.server;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Timeout;
import reactor.netty.DisposableServer;
import reactor.netty.http.server.HttpServer;
import org.springframework.ai.mcp.server.transport.SseServerTransport;
import org.springframework.ai.mcp.spec.ServerMcpTransport;
import org.springframework.http.server.reactive.HttpHandler;
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
import org.springframework.web.reactive.function.server.RouterFunctions;
/**
* Tests for {@link McpSyncServer} using {@link SseServerTransport}.
*
* @author Christian Tzolov
*/
@Timeout(15) // Giving extra time beyond the client timeout
class SseMcpSyncServerTests extends AbstractMcpSyncServerTests {
private static final int PORT = 8182;
private static final String MESSAGE_ENDPOINT = "/mcp/message";
private DisposableServer httpServer;
private SseServerTransport transport;
@Override
protected ServerMcpTransport createMcpTransport() {
transport = new SseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
return transport;
}
@Override
protected void onStart() {
HttpHandler httpHandler = RouterFunctions.toHttpHandler(transport.getRouterFunction());
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
}
@Override
protected void onClose() {
if (httpServer != null) {
httpServer.disposeNow();
}
}
}

439
mcp/src/test/java/org/springframework/ai/mcp/server/_SseAsyncIntegrationTests.java_ Normal file
View file

@ -0,0 +1,439 @@
/*
* Copyright 2024 - 2024 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.
*/
package org.springframework.ai.mcp.server;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import reactor.netty.DisposableServer;
import reactor.netty.http.server.HttpServer;
import reactor.test.StepVerifier;
import org.springframework.ai.mcp.client.McpClient;
import org.springframework.ai.mcp.client.transport.SseClientTransport;
import org.springframework.ai.mcp.server.McpServer.ToolRegistration;
import org.springframework.ai.mcp.server.transport.SseServerTransport;
import org.springframework.ai.mcp.spec.McpError;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
import org.springframework.ai.mcp.spec.McpSchema.CreateMessageRequest;
import org.springframework.ai.mcp.spec.McpSchema.CreateMessageResult;
import org.springframework.ai.mcp.spec.McpSchema.InitializeResult;
import org.springframework.ai.mcp.spec.McpSchema.Role;
import org.springframework.ai.mcp.spec.McpSchema.Root;
import org.springframework.ai.mcp.spec.McpSchema.ServerCapabilities;
import org.springframework.ai.mcp.spec.McpSchema.Tool;
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 static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.awaitility.Awaitility.await;
public class SseAsyncIntegrationTests {
private static final int PORT = 8181;
private static final String MESSAGE_ENDPOINT = "/mcp/message";
private DisposableServer httpServer;
private SseServerTransport mcpServerTransport;
McpClient.Builder clientBuilder;
@BeforeEach
public void before() {
this.mcpServerTransport = new SseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpServerTransport.getRouterFunction());
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
this.httpServer = HttpServer.create().port(PORT).handle(adapter).bindNow();
this.clientBuilder = McpClient
.using(new SseClientTransport(WebClient.builder().baseUrl("http://localhost:" + PORT)));
}
@AfterEach
public void after() {
if (httpServer != null) {
httpServer.disposeNow();
}
}
// ---------------------------------------
// Sampling Tests
// ---------------------------------------
@Test
void testCreateMessageWithoutInitialization() {
var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async();
var messages = List
.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message")));
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of());
StepVerifier.create(mcpAsyncServer.createMessage(request)).verifyErrorSatisfies(error -> {
assertThat(error).isInstanceOf(McpError.class)
.hasMessage("Client must be initialized. Call the initialize method first!");
});
}
@Test
void testCreateMessageWithoutSamplingCapabilities() {
var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async();
var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).sync();
InitializeResult initResult = client.initialize();
assertThat(initResult).isNotNull();
var messages = List
.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message")));
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of());
StepVerifier.create(mcpAsyncServer.createMessage(request)).verifyErrorSatisfies(error -> {
assertThat(error).isInstanceOf(McpError.class)
.hasMessage("Client must be configured with sampling capabilities");
});
}
@Test
void testCreateMessageSuccess() throws InterruptedException {
var mcpAsyncServer = McpServer.using(mcpServerTransport).serverInfo("test-server", "1.0.0").async();
Function<CreateMessageRequest, CreateMessageResult> 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);
};
var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0"))
.capabilities(ClientCapabilities.builder().sampling().build())
.sampling(samplingHandler)
.sync();
InitializeResult initResult = client.initialize();
assertThat(initResult).isNotNull();
var messages = List
.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message")));
var modelPrefs = new McpSchema.ModelPreferences(List.of(), 1.0, 1.0, 1.0);
var request = new McpSchema.CreateMessageRequest(messages, modelPrefs, null,
McpSchema.CreateMessageRequest.ContextInclusionStrategy.NONE, null, 100, List.of(), Map.of());
StepVerifier.create(mcpAsyncServer.createMessage(request)).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();
}
// ---------------------------------------
// Roots Tests
// ---------------------------------------
@Test
void testRootsSuccess() {
List<Root> roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2"));
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
var mcpServer = McpServer.using(mcpServerTransport)
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
.sync();
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
.roots(roots)
.sync();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
assertThat(rootsRef.get()).isNull();
assertThat(mcpServer.listRoots().roots()).containsAll(roots);
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));
});
mcpClient.close();
mcpServer.close();
}
@Test
void testRootsWithoutCapability() {
var mcpServer = McpServer.using(mcpServerTransport).rootsChangeConsumer(rootsUpdate -> {
}).sync();
// Create client without roots capability
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()) // No
// roots
// capability
.sync();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
// Attempt to list roots should fail
assertThatThrownBy(() -> mcpServer.listRoots().roots()).isInstanceOf(McpError.class)
.hasMessage("Roots not supported");
mcpClient.close();
mcpServer.close();
}
@Test
void testRootsWithEmptyRootsList() {
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
var mcpServer = McpServer.using(mcpServerTransport)
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
.sync();
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
.roots(List.of()) // Empty roots list
.sync();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
mcpClient.rootsListChangedNotification();
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
assertThat(rootsRef.get()).isEmpty();
});
mcpClient.close();
mcpServer.close();
}
@Test
void testRootsWithMultipleConsumers() {
List<Root> roots = List.of(new Root("uri1://", "root1"));
AtomicReference<List<Root>> rootsRef1 = new AtomicReference<>();
AtomicReference<List<Root>> rootsRef2 = new AtomicReference<>();
var mcpServer = McpServer.using(mcpServerTransport)
.rootsChangeConsumer(rootsUpdate -> rootsRef1.set(rootsUpdate))
.rootsChangeConsumer(rootsUpdate -> rootsRef2.set(rootsUpdate))
.sync();
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
.roots(roots)
.sync();
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);
});
mcpClient.close();
mcpServer.close();
}
@Test
void testRootsServerCloseWithActiveSubscription() {
List<Root> roots = List.of(new Root("uri1://", "root1"));
AtomicReference<List<Root>> rootsRef = new AtomicReference<>();
var mcpServer = McpServer.using(mcpServerTransport)
.rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate))
.sync();
var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build())
.roots(roots)
.sync();
InitializeResult initResult = mcpClient.initialize();
assertThat(initResult).isNotNull();
mcpClient.rootsListChangedNotification();
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
assertThat(rootsRef.get()).containsAll(roots);
});
// Close server while subscription is active
mcpServer.close();
// Verify client can handle server closure gracefully
mcpClient.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);
ToolRegistration tool1 = new ToolRegistration(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema),
request -> {
// perform a blocking call to a remote service
String response = RestClient.create()
.get()
.uri("https://github.com/spring-projects-experimental/spring-ai-mcp/blob/main/README.md")
.retrieve()
.body(String.class);
assertThat(response).isNotBlank();
return callResponse;
});
var mcpServer = McpServer.using(mcpServerTransport)
.capabilities(ServerCapabilities.builder().tools(true).build())
.tools(tool1)
.sync();
var mcpClient = clientBuilder.sync();
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);
mcpClient.close();
mcpServer.close();
}
@Test
void testToolListChangeHandlingSuccess() {
var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
ToolRegistration tool1 = new ToolRegistration(new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema),
request -> {
// perform a blocking call to a remote service
String response = RestClient.create()
.get()
.uri("https://github.com/spring-projects-experimental/spring-ai-mcp/blob/main/README.md")
.retrieve()
.body(String.class);
assertThat(response).isNotBlank();
return callResponse;
});
var mcpServer = McpServer.using(mcpServerTransport)
.capabilities(ServerCapabilities.builder().tools(true).build())
.tools(tool1)
.sync();
AtomicReference<List<Tool>> rootsRef = new AtomicReference<>();
var mcpClient = clientBuilder.toolsChangeConsumer(toolsUpdate -> {
// perform a blocking call to a remote service
String response = RestClient.create()
.get()
.uri("https://github.com/spring-projects-experimental/spring-ai-mcp/blob/main/README.md")
.retrieve()
.body(String.class);
assertThat(response).isNotBlank();
rootsRef.set(toolsUpdate);
}).sync();
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
ToolRegistration tool2 = new ToolRegistration(new McpSchema.Tool("tool2", "tool2 description", emptyJsonSchema),
request -> callResponse);
mcpServer.addTool(tool2);
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
assertThat(rootsRef.get()).containsAll(List.of(tool2.tool()));
});
mcpClient.close();
mcpServer.close();
}
}

283
mcp/src/test/java/org/springframework/ai/mcp/server/transport/SseServerTransportTests.java
View file

@ -1,283 +0,0 @@
package org.springframework.ai.mcp.server.transport;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import com.fasterxml.jackson.databind.ObjectMapper;
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.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.test.StepVerifier;
import org.springframework.ai.mcp.spec.McpSchema;
import org.springframework.ai.mcp.spec.McpSchema.JSONRPCRequest;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.test.web.reactive.server.WebTestClient;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/**
* Tests for the {@link SseServerTransport} class.
*
* @author Christian Tzolov
*/
@Timeout(15)
class SseServerTransportTests {
private static final Logger logger = LoggerFactory.getLogger(SseServerTransportTests.class);
private ObjectMapper objectMapper;
private String messageEndpoint;
private SseServerTransport transport;
private WebTestClient webTestClient;
@BeforeEach
void setUp() {
objectMapper = new ObjectMapper();
messageEndpoint = "/message";
transport = new SseServerTransport(objectMapper, messageEndpoint);
webTestClient = WebTestClient.bindToRouterFunction(transport.getRouterFunction()).build();
}
@Test
void constructorValidation() {
assertThatThrownBy(() -> new SseServerTransport(null, "/message")).isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("ObjectMapper must not be null");
assertThatThrownBy(() -> new SseServerTransport(new ObjectMapper(), null))
.isInstanceOf(IllegalArgumentException.class)
.hasMessageContaining("Message endpoint must not be null");
}
@Test
void testSseConnectionEstablishment() {
List<ServerSentEvent<String>> events = new ArrayList<>();
webTestClient.get()
.uri("/sse")
.accept(MediaType.TEXT_EVENT_STREAM)
.exchange()
.expectStatus()
.isOk()
.expectHeader()
.contentTypeCompatibleWith(MediaType.TEXT_EVENT_STREAM)
.returnResult(String.class)
.getResponseBody()
.map(data -> ServerSentEvent.<String>builder().data(data).build())
.take(1) // Take only the initial endpoint event
.subscribe(events::add);
// Wait a bit for the event to be received
StepVerifier.create(Mono.delay(Duration.ofMillis(500))).expectNextCount(1).verifyComplete();
assertThat(events).hasSize(1);
assertThat(events.get(0).data()).isEqualTo(messageEndpoint);
}
@Test
void testMessageHandling() {
AtomicReference<McpSchema.JSONRPCMessage> receivedMessage = new AtomicReference<>();
// Set up message handler
transport.connect(message -> {
message.doOnNext(receivedMessage::set).subscribe();
return Mono.empty();
}).block();
// Create a test message
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test-method", "test-id",
Map.of("key", "value"));
// Send message to endpoint
webTestClient.post()
.uri(messageEndpoint)
.contentType(MediaType.APPLICATION_JSON)
.bodyValue(testMessage)
.exchange()
.expectStatus()
.isOk();
// Verify the message was received and processed
assertThat(receivedMessage.get()).isNotNull();
McpSchema.JSONRPCRequest receivedRequest = (McpSchema.JSONRPCRequest) receivedMessage.get();
assertThat(receivedRequest.id()).isEqualTo(testMessage.id());
assertThat(receivedRequest.method()).isEqualTo(testMessage.method());
}
@Test
@Disabled("Flaky test")
void testBroadcastMessage() {
// Create test clients
int clientCount = 3;
CountDownLatch connectLatch = new CountDownLatch(clientCount);
CountDownLatch messageLatch = new CountDownLatch(clientCount);
List<List<ServerSentEvent<String>>> allReceivedEvents = new ArrayList<>();
List<Disposable> subscriptions = new ArrayList<>();
// Connect clients
for (int i = 0; i < clientCount; i++) {
List<ServerSentEvent<String>> clientEvents = new ArrayList<>();
allReceivedEvents.add(clientEvents);
Flux<ServerSentEvent<String>> eventStream = webTestClient.get()
.uri("/sse")
.accept(MediaType.TEXT_EVENT_STREAM)
.exchange()
.expectStatus()
.isOk()
.returnResult(String.class)
.getResponseBody()
.map(data -> ServerSentEvent.<String>builder().data(data).build());
Disposable subscription = eventStream.doOnNext(event -> {
clientEvents.add(event);
if (event.event() != null && event.event().equals(SseServerTransport.ENDPOINT_EVENT_TYPE)) {
connectLatch.countDown();
}
else if (event.event() != null && event.event().equals(SseServerTransport.MESSAGE_EVENT_TYPE)) {
messageLatch.countDown();
}
}).subscribe();
subscriptions.add(subscription);
}
// Wait for all clients to connect
try {
assertThat(connectLatch.await(5, TimeUnit.SECONDS)).isTrue();
}
catch (InterruptedException e) {
throw new RuntimeException(e);
}
// Verify initial connections
for (List<ServerSentEvent<String>> events : allReceivedEvents) {
assertThat(events).hasSize(1);
assertThat(events.get(0).data()).isEqualTo(messageEndpoint);
}
// Give clients time to fully establish their subscriptions
logger.debug("Waiting for subscriptions to stabilize...");
try {
Thread.sleep(1000);
}
catch (InterruptedException e) {
throw new RuntimeException(e);
}
logger.debug("Sending broadcast message to {} clients", clientCount);
// Send broadcast message
JSONRPCRequest broadcastMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "broadcast", "broadcast-id",
Map.of("message", "Hello all!"));
// Send the message
transport.sendMessage(broadcastMessage).block(Duration.ofSeconds(5));
// Wait for all clients to receive the broadcast
try {
assertThat(messageLatch.await(5, TimeUnit.SECONDS)).isTrue();
}
catch (InterruptedException e) {
throw new RuntimeException(e);
}
// Verify each client received both messages
for (List<ServerSentEvent<String>> events : allReceivedEvents) {
assertThat(events).hasSize(2);
assertThat(events.get(0).data()).isEqualTo(messageEndpoint);
assertThat(events.get(1).data()).contains("broadcast-id");
}
// Cleanup
subscriptions.forEach(Disposable::dispose);
}
@Test
void testGracefulShutdown() {
// Connect a client
Flux<ServerSentEvent<String>> eventStream = webTestClient.get()
.uri("/sse")
.accept(MediaType.TEXT_EVENT_STREAM)
.exchange()
.expectStatus()
.isOk()
.returnResult(String.class)
.getResponseBody()
.map(data -> ServerSentEvent.<String>builder().data(data).build());
List<ServerSentEvent<String>> receivedEvents = new ArrayList<>();
eventStream.subscribe(receivedEvents::add);
// Wait for connection
StepVerifier.create(Mono.delay(Duration.ofMillis(500))).expectNextCount(1).verifyComplete();
// Initiate shutdown
transport.closeGracefully().block(Duration.ofSeconds(5));
// Verify server rejects new connections with timeout
webTestClient.get()
.uri("/sse")
.accept(MediaType.TEXT_EVENT_STREAM)
.exchange()
.expectStatus()
.isEqualTo(HttpStatus.SERVICE_UNAVAILABLE)
.expectBody(String.class)
.isEqualTo("Server is shutting down");
// Verify server rejects new messages with timeout
webTestClient.post()
.uri(messageEndpoint)
.contentType(MediaType.APPLICATION_JSON)
.bodyValue("""
{
"jsonrpc": "2.0",
"method": "test",
"id": "1",
"params": {}
}
""")
.exchange()
.expectStatus()
.isEqualTo(HttpStatus.SERVICE_UNAVAILABLE)
.expectBody(String.class)
.isEqualTo("Server is shutting down");
}
@Test
void testInvalidMessageHandling() {
// Test invalid JSON
webTestClient.post()
.uri(messageEndpoint)
.contentType(MediaType.APPLICATION_JSON)
.bodyValue("invalid json")
.exchange()
.expectStatus()
.isBadRequest();
// Test invalid message format
webTestClient.post().uri(messageEndpoint).contentType(MediaType.APPLICATION_JSON).bodyValue("""
{
"invalid": "message"
}
""").exchange().expectStatus().isBadRequest();
}
}