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(transport): Simplify WebMvc SSE transport with native Spring SseBuilder

Browse source 
Replaces custom BlockingQueue-based SSE implementation with Spring's native
SseBuilder for simpler and more reliable event handling.

This change:
- Removes custom SSEEvent and queue implementation
- Uses Spring's SseBuilder directly for SSE event management
- Improves error handling with direct error reporting
- Reduces complexity in session management
This commit is contained in:
Christian Tzolov 2025-01-11 18:25:17 +01:00
commit bfd07b4cda
1 changed files with 18 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

47
mcp-transport/mcp-webmvc-sse-transport/src/main/java/org/springframework/ai/mcp/server/transport/WebMvcSseServerTransport.java
View file

@ -18,10 +18,7 @@ package org.springframework.ai.mcp.server.transport;
import java.io.IOException; import java.io.IOException;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.function.Function; import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.core.type.TypeReference;
@ -39,6 +36,7 @@ import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.RouterFunctions; import org.springframework.web.servlet.function.RouterFunctions;
import org.springframework.web.servlet.function.ServerRequest; import org.springframework.web.servlet.function.ServerRequest;
import org.springframework.web.servlet.function.ServerResponse; import org.springframework.web.servlet.function.ServerResponse;
import org.springframework.web.servlet.function.ServerResponse.SseBuilder;
/** /**
* Server-side implementation of the MCP HTTP with SSE transport specification using * Server-side implementation of the MCP HTTP with SSE transport specification using
@ -122,11 +120,11 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
sessions.values().forEach(session -> { sessions.values().forEach(session -> {
try { try {
var event = new SSEEvent(session.id, MESSAGE_EVENT_TYPE, jsonText); session.sseBuilder.id(session.id).event(MESSAGE_EVENT_TYPE).data(jsonText);
session.queue.put(event);
} }
catch (Exception e) { catch (Exception e) {
logger.error("Failed to send message to session {}: {}", session.id, e.getMessage()); logger.error("Failed to send message to session {}: {}", session.id, e.getMessage());
session.sseBuilder.error(e);
} }
}); });
} }
@ -147,27 +145,20 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
String sessionId = UUID.randomUUID().toString(); String sessionId = UUID.randomUUID().toString();
logger.debug("Creating new SSE connection for session: {}", sessionId); logger.debug("Creating new SSE connection for session: {}", sessionId);
ClientSession session = new ClientSession(sessionId);
this.sessions.put(sessionId, session);
// Send initial endpoint event // Send initial endpoint event
try { try {
session.queue.put(new SSEEvent(session.id, ENDPOINT_EVENT_TYPE, messageEndpoint));
return ServerResponse.sse(sseBuilder -> { return ServerResponse.sse(sseBuilder -> {
// new Thread(() -> {
while (!this.isClosing) { ClientSession session = new ClientSession(sessionId, sseBuilder);
try { this.sessions.put(sessionId, session);
SSEEvent sseEvent = session.queue.poll(100, TimeUnit.MILLISECONDS);
if (sseEvent != null) { try {
sseBuilder.id(sseEvent.id).event(sseEvent.type()).data(sseEvent.data()); session.sseBuilder.id(session.id).event(ENDPOINT_EVENT_TYPE).data(messageEndpoint);
} }
} catch (Exception e) {
catch (Exception e) { logger.error("Failed to poll event from session queue: {}", e.getMessage());
logger.error("Failed to poll event from session queue: {}", e.getMessage()); sseBuilder.error(e);
sseBuilder.error(e);
}
} }
// }).start();
}); });
} }
catch (Exception e) { catch (Exception e) {
@ -205,9 +196,6 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
} }
} }
record SSEEvent(String id, String type, String data) {
}
/** /**
* Represents an active client session. * Represents an active client session.
*/ */
@ -215,22 +203,23 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
private final String id; private final String id;
private final BlockingQueue<SSEEvent> queue; private final SseBuilder sseBuilder;
ClientSession(String id) { ClientSession(String id, SseBuilder sseBuilder) {
this.id = id; this.id = id;
this.queue = new LinkedBlockingQueue<>(); this.sseBuilder = sseBuilder;
logger.debug("Session {} initialized with SSE emitter", id); logger.debug("Session {} initialized with SSE emitter", id);
} }
void close() { void close() {
logger.debug("Closing session: {}", id); logger.debug("Closing session: {}", id);
try { try {
queue.remove(); sseBuilder.complete();
logger.debug("Successfully completed SSE emitter for session {}", id); logger.debug("Successfully completed SSE emitter for session {}", id);
} }
catch (Exception e) { catch (Exception e) {
logger.warn("Failed to complete SSE emitter for session {}: {}", id, e.getMessage()); logger.warn("Failed to complete SSE emitter for session {}: {}", id, e.getMessage());
// sseBuilder.error(e);
} }
} }