refactor: Add HttpServletSseServerTransport for generic Servlet integration
Introduces a new Server-Sent Events (SSE) Server transport implementation based on the Servlet API that enables MCP integration with any Java HTTP server application supporting Servlets. This implementation complements the existing WebFlux and WebMvc transports by providing a generic Servlet-based option that works with any Java web server, not just Spring applications. Key changes: - Add HttpServletSseServerTransport with async Servlet 6.0 support - Add integration tests with embedded Tomcat - Add servlet dependencies with externalized versions: - jakarta.servlet-api 6.1.0 (provided) - tomcat-embed-core 11.0.2 (test) - tomcat-embed-websocket 11.0.2 (test) Resolves #61
This commit is contained in:
Christian Tzolov
parent
42b5931dd5
commit
0c8968044c
11 changed files with 273 additions and 206 deletions
|
|
@ -153,6 +153,10 @@ class StdioServerTransport implements ServerMcpTransport {
|
|||
}
|
||||
|
||||
|
||||
class HttpServletSseServerTransport implements ServerMcpTransport {
|
||||
}
|
||||
|
||||
|
||||
class HttpClientSseClientTransport implements ClientMcpTransport {
|
||||
}
|
||||
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
Before Width: | Height: | Size: 65 KiB After Width: | Height: | Size: 67 KiB |
|
|
@ -15,7 +15,8 @@ Java SDK implementation of the link:https://modelcontextprotocol.io/introduction
|
|||
* Multiple transport implementations:
|
||||
** Core transports:
|
||||
*** Stdio-based transport for process-based communication
|
||||
*** Java HttpClient-based SSE transport for HTTP streaming
|
||||
*** Java HttpClient-based SSE client transport for HTTP SSE Client-side streaming
|
||||
*** Servlet-based SSE server transport for HTTP SSE Server streaming
|
||||
** Spring-based transports:
|
||||
*** WebFlux SSE transport for reactive HTTP streaming
|
||||
*** WebMVC SSE transport for servlet-based HTTP streaming
|
||||
|
|
@ -515,6 +516,41 @@ Implements the MCP HTTP with SSE transport specification, providing:
|
|||
* Message routing and session management
|
||||
* Graceful shutdown capabilities
|
||||
|
||||
SSE (Servlet)::
|
||||
+
|
||||
Creates a Servlet-based SSE server transport.
|
||||
Included in the core `mcp` module.
|
||||
The `HttpServletSseServerTransport` can be used with any Servlet container.
|
||||
To using it with a Spring Web application, you can register it as a Servlet bean:
|
||||
+
|
||||
[source,java]
|
||||
----
|
||||
@Configuration
|
||||
@EnableWebMvc
|
||||
public class McpServerConfig implements WebMvcConfigurer {
|
||||
|
||||
@Bean
|
||||
public HttpServletSseServerTransport servletSseServerTransport() {
|
||||
return new HttpServletSseServerTransport(new ObjectMapper(), "/mcp/message");
|
||||
}
|
||||
|
||||
@Bean
|
||||
public ServletRegistrationBean customServletBean(HttpServletSseServerTransport servlet) {
|
||||
return new ServletRegistrationBean(servlet);
|
||||
}
|
||||
}
|
||||
----
|
||||
+
|
||||
Implements the MCP HTTP with SSE transport specification using the traditional Servlet API, providing:
|
||||
+
|
||||
* Asynchronous message handling using Servlet 6.0 async support
|
||||
* Session management for multiple client connections
|
||||
* Two types of endpoints:
|
||||
** SSE endpoint (/sse) for server-to-client events
|
||||
** Message endpoint (configurable) for client-to-server requests
|
||||
* Error handling and response formatting
|
||||
* Graceful shutdown support
|
||||
|
||||
======
|
||||
|
||||
=== Server Capabilities
|
||||
|
|
|
|||
|
|
@ -20,7 +20,8 @@ Core implementation of the Model Context Protocol specification, providing:
|
|||
|
||||
* *Core Transports*
|
||||
** Stdio-based (`StdioClientTransport`, `StdioServerTransport`) for process-based communication
|
||||
** Java HttpClient-based SSE client (`HttpClientSseClientTransport`) for HTTP streaming
|
||||
** Java HttpClient-based SSE client (`HttpClientSseClientTransport`) for HTTP SSE Client-side streaming
|
||||
** Servlet-based SSE server (`HttpServletSseServerTransport`) for HTTP SSE Server streaming using traditional Servlet API
|
||||
|
||||
* *Optional SSE Transports*
|
||||
** link:https://github.com/spring-projects-experimental/spring-ai-mcp/tree/main/mcp-transport/mcp-webflux-sse-transport[WebFlux SSE Transport] - Reactive HTTP streaming with Spring WebFlux (Client & Server)
|
||||
|
|
@ -114,6 +115,7 @@ Reffer to the xref:dependency-management.adoc[Dependency Management] page for mo
|
|||
* link:https://github.com/habuma/spring-ai-examples/tree/main/spring-ai-mcp[Theme Park API] - MCP server/client with Theme Park API tools
|
||||
* link:https://github.com/spring-projects/spring-ai-examples/tree/main/model-context-protocol/mcp-webflux-server[WebFlux SSE] - WebFlux server/client implementation
|
||||
* link:https://github.com/spring-projects/spring-ai-examples/tree/main/model-context-protocol/mcp-webmvc-server[WebMvc SSE] - WebMvc server with HttpClient implementation
|
||||
* link:https://github.com/spring-projects/spring-ai-examples/tree/main/model-context-protocol/mcp-servlet-server[Servlet SSE] - SSE Servlet server with HttpClient implementation
|
||||
|
||||
== Documentation
|
||||
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@
|
|||
<dependency>
|
||||
<groupId>org.awaitility</groupId>
|
||||
<artifactId>awaitility</artifactId>
|
||||
<version>4.2.0</version>
|
||||
<version>${awaitility.version}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
|
|
@ -80,21 +80,6 @@
|
|||
<artifactId>logback-classic</artifactId>
|
||||
<version>${logback.version}</version>
|
||||
</dependency>
|
||||
|
||||
<!-- <dependency>
|
||||
<groupId>jakarta.servlet</groupId>
|
||||
<artifactId>jakarta.servlet-api</artifactId>
|
||||
<version>5.0.0</version>
|
||||
<scope>provided</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.apache.tomcat.embed</groupId>
|
||||
<artifactId>tomcat-embed-core</artifactId>
|
||||
<version>10.1.18</version>
|
||||
<scope>test</scope>
|
||||
</dependency> -->
|
||||
|
||||
</dependencies>
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -97,7 +97,7 @@
|
|||
<dependency>
|
||||
<groupId>org.awaitility</groupId>
|
||||
<artifactId>awaitility</artifactId>
|
||||
<version>4.2.0</version>
|
||||
<version>${awaitility.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
|
|
|
|||
|
|
@ -87,7 +87,7 @@
|
|||
<dependency>
|
||||
<groupId>org.awaitility</groupId>
|
||||
<artifactId>awaitility</artifactId>
|
||||
<version>4.2.0</version>
|
||||
<version>${awaitility.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
|
|
@ -111,14 +111,14 @@
|
|||
<dependency>
|
||||
<groupId>jakarta.servlet</groupId>
|
||||
<artifactId>jakarta.servlet-api</artifactId>
|
||||
<version>5.0.0</version>
|
||||
<version>${jakarta.servlet.version}</version>
|
||||
<scope>provided</scope>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>org.apache.tomcat.embed</groupId>
|
||||
<artifactId>tomcat-embed-core</artifactId>
|
||||
<version>11.0.2</version>
|
||||
<version>${tomcat.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
|
|
|
|||
26
mcp/pom.xml
26
mcp/pom.xml
|
|
@ -103,7 +103,7 @@
|
|||
<dependency>
|
||||
<groupId>org.awaitility</groupId>
|
||||
<artifactId>awaitility</artifactId>
|
||||
<version>4.2.0</version>
|
||||
<version>${awaitility.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
|
|
@ -114,7 +114,29 @@
|
|||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
<!-- Used by the HttpServletSseServerTransport -->
|
||||
<dependency>
|
||||
<groupId>jakarta.servlet</groupId>
|
||||
<artifactId>jakarta.servlet-api</artifactId>
|
||||
<version>${jakarta.servlet.version}</version>
|
||||
<scope>provided</scope>
|
||||
</dependency>
|
||||
|
||||
<!-- Tomcat dependencies for testing -->
|
||||
<dependency>
|
||||
<groupId>org.apache.tomcat.embed</groupId>
|
||||
<artifactId>tomcat-embed-core</artifactId>
|
||||
<version>${tomcat.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
<dependency>
|
||||
<groupId>org.apache.tomcat.embed</groupId>
|
||||
<artifactId>tomcat-embed-websocket</artifactId>
|
||||
<version>${tomcat.version}</version>
|
||||
<scope>test</scope>
|
||||
</dependency>
|
||||
|
||||
</dependencies>
|
||||
|
||||
|
||||
</project>
|
||||
</project>
|
||||
|
|
|
|||
|
|
@ -0,0 +1,407 @@
|
|||
/*
|
||||
* 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 Server-Sent Events (SSE) transport
|
||||
* specification. This implementation provides similar functionality to
|
||||
* WebFluxSseServerTransport but uses the traditional Servlet API instead of WebFlux.
|
||||
*
|
||||
* <p>
|
||||
* The transport handles two types of endpoints:
|
||||
* <ul>
|
||||
* <li>SSE endpoint (/sse) - Establishes a long-lived connection for server-to-client
|
||||
* events</li>
|
||||
* <li>Message endpoint (configurable) - Handles client-to-server message requests</li>
|
||||
* </ul>
|
||||
*
|
||||
* <p>
|
||||
* Features:
|
||||
* <ul>
|
||||
* <li>Asynchronous message handling using Servlet 6.0 async support</li>
|
||||
* <li>Session management for multiple client connections</li>
|
||||
* <li>Graceful shutdown support</li>
|
||||
* <li>Error handling and response formatting</li>
|
||||
* </ul>
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @see ServerMcpTransport
|
||||
* @see HttpServlet
|
||||
*/
|
||||
|
||||
@WebServlet(asyncSupported = true)
|
||||
public class HttpServletSseServerTransport extends HttpServlet implements ServerMcpTransport {
|
||||
|
||||
/** Logger for this class */
|
||||
private static final Logger logger = LoggerFactory.getLogger(HttpServletSseServerTransport.class);
|
||||
|
||||
/** The endpoint path for SSE connections */
|
||||
public static final String SSE_ENDPOINT = "/sse";
|
||||
|
||||
/** Event type for regular messages */
|
||||
public static final String MESSAGE_EVENT_TYPE = "message";
|
||||
|
||||
/** Event type for endpoint information */
|
||||
public static final String ENDPOINT_EVENT_TYPE = "endpoint";
|
||||
|
||||
/** JSON object mapper for serialization/deserialization */
|
||||
private final ObjectMapper objectMapper;
|
||||
|
||||
/** The endpoint path for handling client messages */
|
||||
private final String messageEndpoint;
|
||||
|
||||
/** Map of active client sessions, keyed by session ID */
|
||||
private final Map<String, ClientSession> sessions = new ConcurrentHashMap<>();
|
||||
|
||||
/** Flag indicating if the transport is in the process of shutting down */
|
||||
private final AtomicBoolean isClosing = new AtomicBoolean(false);
|
||||
|
||||
/** Handler for processing incoming messages */
|
||||
private Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> connectHandler;
|
||||
|
||||
/**
|
||||
* Creates a new HttpServletSseServerTransport instance.
|
||||
* @param objectMapper The JSON object mapper to use for message
|
||||
* serialization/deserialization
|
||||
* @param messageEndpoint The endpoint path where clients will send their messages
|
||||
*/
|
||||
public HttpServletSseServerTransport(ObjectMapper objectMapper, String messageEndpoint) {
|
||||
this.objectMapper = objectMapper;
|
||||
this.messageEndpoint = messageEndpoint;
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles GET requests to establish SSE connections.
|
||||
* <p>
|
||||
* This method sets up a new SSE connection when a client connects to the SSE
|
||||
* endpoint. It configures the response headers for SSE, creates a new session, and
|
||||
* sends the initial endpoint information to the client.
|
||||
* @param request The HTTP servlet request
|
||||
* @param response The HTTP servlet response
|
||||
* @throws ServletException If a servlet-specific error occurs
|
||||
* @throws IOException If an I/O error occurs
|
||||
*/
|
||||
@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);
|
||||
this.sessions.put(sessionId, session);
|
||||
|
||||
// Send initial endpoint event
|
||||
this.sendEvent(writer, ENDPOINT_EVENT_TYPE, messageEndpoint);
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles POST requests for client messages.
|
||||
* <p>
|
||||
* This method processes incoming messages from clients, routes them through the
|
||||
* connect handler if configured, and sends back the appropriate response. It handles
|
||||
* error cases and formats error responses according to the MCP specification.
|
||||
* @param request The HTTP servlet request
|
||||
* @param response The HTTP servlet response
|
||||
* @throws ServletException If a servlet-specific error occurs
|
||||
* @throws IOException If an I/O error occurs
|
||||
*/
|
||||
@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");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Sets up the message handler for processing client requests.
|
||||
* @param handler The function to process incoming messages and produce responses
|
||||
* @return A Mono that completes when the handler is set up
|
||||
*/
|
||||
@Override
|
||||
public Mono<Void> connect(Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
|
||||
this.connectHandler = handler;
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
/**
|
||||
* Broadcasts a message to all connected clients.
|
||||
* <p>
|
||||
* This method serializes the message and sends it to all active client sessions. If a
|
||||
* client is disconnected, its session is removed.
|
||||
* @param message The message to broadcast
|
||||
* @return A Mono that completes when the message has been sent to all clients
|
||||
*/
|
||||
@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 {
|
||||
this.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()));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Closes the transport.
|
||||
* <p>
|
||||
* This implementation delegates to the super class's close method.
|
||||
*/
|
||||
@Override
|
||||
public void close() {
|
||||
ServerMcpTransport.super.close();
|
||||
}
|
||||
|
||||
/**
|
||||
* Unmarshals data from one type to another using the object mapper.
|
||||
* @param <T> The target type
|
||||
* @param data The source data
|
||||
* @param typeRef The type reference for the target type
|
||||
* @return The unmarshaled data
|
||||
*/
|
||||
@Override
|
||||
public <T> T unmarshalFrom(Object data, TypeReference<T> typeRef) {
|
||||
return objectMapper.convertValue(data, typeRef);
|
||||
}
|
||||
|
||||
/**
|
||||
* Initiates a graceful shutdown of the transport.
|
||||
* <p>
|
||||
* This method marks the transport as closing and closes all active client sessions.
|
||||
* New connection attempts will be rejected during shutdown.
|
||||
* @return A Mono that completes when all sessions have been closed
|
||||
*/
|
||||
@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();
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Sends an SSE event to a client.
|
||||
* @param writer The writer to send the event through
|
||||
* @param eventType The type of event (message or endpoint)
|
||||
* @param data The event data
|
||||
* @throws IOException If an error occurs while writing the event
|
||||
*/
|
||||
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");
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Removes a client session and completes its async context.
|
||||
* @param session The session to remove
|
||||
*/
|
||||
private void removeSession(ClientSession session) {
|
||||
sessions.remove(session.id);
|
||||
session.asyncContext.complete();
|
||||
}
|
||||
|
||||
/**
|
||||
* Represents a client connection session.
|
||||
* <p>
|
||||
* This class holds the necessary information about a client's SSE connection,
|
||||
* including its ID, async context, and output writer.
|
||||
*/
|
||||
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;
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
/**
|
||||
* Cleans up resources when the servlet is being destroyed.
|
||||
* <p>
|
||||
* This method ensures a graceful shutdown by closing all client connections before
|
||||
* calling the parent's destroy method.
|
||||
*/
|
||||
@Override
|
||||
public void destroy() {
|
||||
closeGracefully().block();
|
||||
super.destroy();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -0,0 +1,344 @@
|
|||
/*
|
||||
* 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.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.apache.catalina.Context;
|
||||
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.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import org.springframework.ai.mcp.client.McpClient;
|
||||
import org.springframework.ai.mcp.client.transport.HttpClientSseClientTransport;
|
||||
import org.springframework.ai.mcp.server.McpServer;
|
||||
import org.springframework.ai.mcp.server.McpServer.ToolRegistration;
|
||||
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.web.client.RestClient;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.awaitility.Awaitility.await;
|
||||
|
||||
public class HttpServletSseServerTransportIntegrationTests {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(HttpServletSseServerTransportIntegrationTests.class);
|
||||
|
||||
private static final int PORT = 8184;
|
||||
|
||||
private static final String MESSAGE_ENDPOINT = "/mcp/message";
|
||||
|
||||
private HttpServletSseServerTransport mcpServerTransport;
|
||||
|
||||
McpClient.Builder clientBuilder;
|
||||
|
||||
private Tomcat tomcat;
|
||||
|
||||
@BeforeEach
|
||||
public void before() {
|
||||
tomcat = new Tomcat();
|
||||
tomcat.setPort(PORT);
|
||||
|
||||
String baseDir = System.getProperty("java.io.tmpdir");
|
||||
tomcat.setBaseDir(baseDir);
|
||||
|
||||
Context context = tomcat.addContext("", baseDir);
|
||||
|
||||
// Create and configure the transport
|
||||
mcpServerTransport = new HttpServletSseServerTransport(new ObjectMapper(), MESSAGE_ENDPOINT);
|
||||
|
||||
// Add transport servlet to Tomcat
|
||||
org.apache.catalina.Wrapper wrapper = context.createWrapper();
|
||||
wrapper.setName("mcpServlet");
|
||||
wrapper.setServlet(mcpServerTransport);
|
||||
wrapper.setLoadOnStartup(1);
|
||||
wrapper.setAsyncSupported(true);
|
||||
context.addChild(wrapper);
|
||||
context.addServletMappingDecoded("/*", "mcpServlet");
|
||||
|
||||
try {
|
||||
var connector = tomcat.getConnector();
|
||||
connector.setAsyncTimeout(3000);
|
||||
tomcat.start();
|
||||
assertThat(tomcat.getServer().getState() == LifecycleState.STARTED);
|
||||
}
|
||||
catch (Exception e) {
|
||||
throw new RuntimeException("Failed to start Tomcat", e);
|
||||
}
|
||||
|
||||
this.clientBuilder = McpClient.using(new HttpClientSseClientTransport("http://localhost:" + PORT));
|
||||
}
|
||||
|
||||
@AfterEach
|
||||
public void after() {
|
||||
if (mcpServerTransport != null) {
|
||||
mcpServerTransport.closeGracefully().block();
|
||||
}
|
||||
if (tomcat != null) {
|
||||
try {
|
||||
tomcat.stop();
|
||||
tomcat.destroy();
|
||||
}
|
||||
catch (LifecycleException e) {
|
||||
throw new RuntimeException("Failed to stop Tomcat", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@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() {
|
||||
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();
|
||||
}
|
||||
|
||||
@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);
|
||||
});
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
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 -> {
|
||||
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 -> {
|
||||
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>> toolsRef = new AtomicReference<>();
|
||||
var mcpClient = clientBuilder.toolsChangeConsumer(toolsUpdate -> {
|
||||
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();
|
||||
toolsRef.set(toolsUpdate);
|
||||
}).sync();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
assertThat(toolsRef.get()).isNull();
|
||||
|
||||
assertThat(mcpClient.listTools().tools()).contains(tool1.tool());
|
||||
|
||||
mcpServer.notifyToolsListChanged();
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(toolsRef.get()).containsAll(List.of(tool1.tool()));
|
||||
});
|
||||
|
||||
mcpServer.removeTool("tool1");
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(toolsRef.get()).isEmpty();
|
||||
});
|
||||
|
||||
ToolRegistration tool2 = new ToolRegistration(new McpSchema.Tool("tool2", "tool2 description", emptyJsonSchema),
|
||||
request -> callResponse);
|
||||
|
||||
mcpServer.addTool(tool2);
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(toolsRef.get()).containsAll(List.of(tool2.tool()));
|
||||
});
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testInitialize() {
|
||||
var mcpServer = McpServer.using(mcpServerTransport).sync();
|
||||
var mcpClient = clientBuilder.sync();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
}
|
||||
3
pom.xml
3
pom.xml
|
|
@ -97,6 +97,9 @@
|
|||
<io.spring.maven.antora-version>0.0.4</io.spring.maven.antora-version>
|
||||
<asciidoctorj-pdf.version>1.6.2</asciidoctorj-pdf.version>
|
||||
<junit-jupiter.version>5.10.5</junit-jupiter.version>
|
||||
<tomcat.version>11.0.2</tomcat.version>
|
||||
<jakarta.servlet.version>6.1.0</jakarta.servlet.version>
|
||||
<awaitility.version>4.2.0</awaitility.version>
|
||||
</properties>
|
||||
|
||||
<modules>
|
||||
|
|
|
|||
Loading…
Reference in a new issue