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.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
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.ServerRequest;
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
@ -122,11 +120,11 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
sessions.values().forEach(session -> {
try {
var event = new SSEEvent(session.id, MESSAGE_EVENT_TYPE, jsonText);
session.queue.put(event);
session.sseBuilder.id(session.id).event(MESSAGE_EVENT_TYPE).data(jsonText);
}
catch (Exception e) {
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();
logger.debug("Creating new SSE connection for session: {}", sessionId);
ClientSession session = new ClientSession(sessionId);
this.sessions.put(sessionId, session);
// Send initial endpoint event
try {
session.queue.put(new SSEEvent(session.id, ENDPOINT_EVENT_TYPE, messageEndpoint));
return ServerResponse.sse(sseBuilder -> {
// new Thread(() -> {
while (!this.isClosing) {
try {
SSEEvent sseEvent = session.queue.poll(100, TimeUnit.MILLISECONDS);
if (sseEvent != null) {
sseBuilder.id(sseEvent.id).event(sseEvent.type()).data(sseEvent.data());
}
}
catch (Exception e) {
logger.error("Failed to poll event from session queue: {}", e.getMessage());
sseBuilder.error(e);
}
ClientSession session = new ClientSession(sessionId, sseBuilder);
this.sessions.put(sessionId, session);
try {
session.sseBuilder.id(session.id).event(ENDPOINT_EVENT_TYPE).data(messageEndpoint);
}
catch (Exception e) {
logger.error("Failed to poll event from session queue: {}", e.getMessage());
sseBuilder.error(e);
}
// }).start();
});
}
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.
*/
@ -215,22 +203,23 @@ public class WebMvcSseServerTransport implements ServerMcpTransport {
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.queue = new LinkedBlockingQueue<>();
this.sseBuilder = sseBuilder;
logger.debug("Session {} initialized with SSE emitter", id);
}
void close() {
logger.debug("Closing session: {}", id);
try {
queue.remove();
sseBuilder.complete();
logger.debug("Successfully completed SSE emitter for session {}", id);
}
catch (Exception e) {
logger.warn("Failed to complete SSE emitter for session {}: {}", id, e.getMessage());
// sseBuilder.error(e);
}
}