fix: Improve MCP server stability and test reliability

- Replace bounded elastic schedulers with single thread executors
- Add timeouts to MCP server tests
- Enhance error handling in StdioServerTransport
- Comment out unused PaginatedRequest unmarshalling
- Simplify shutdown logic in StdioServerTransport
This commit is contained in:
Christian Tzolov 2024-12-30 12:15:39 +01:00
commit ad3cdd20a9
9 changed files with 50 additions and 41 deletions

View file

@ -349,9 +349,9 @@ public class McpAsyncServer {
private DefaultMcpSession.RequestHandler resourcesListRequestHandler() { private DefaultMcpSession.RequestHandler resourcesListRequestHandler() {
return params -> { return params -> {
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, // McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
new TypeReference<McpSchema.PaginatedRequest>() { // new TypeReference<McpSchema.PaginatedRequest>() {
}); // });
var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList(); var resourceList = this.resources.values().stream().map(ResourceRegistration::resource).toList();
@ -361,9 +361,9 @@ public class McpAsyncServer {
private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() { private DefaultMcpSession.RequestHandler resourceTemplateListRequestHandler() {
return params -> { return params -> {
McpSchema.PaginatedRequest request = transport.unmarshalFrom(params, // McpSchema.PaginatedRequest request = transport.unmarshalFrom(params,
new TypeReference<McpSchema.PaginatedRequest>() { // new TypeReference<McpSchema.PaginatedRequest>() {
}); // });
return Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null)); return Mono.just(new McpSchema.ListResourceTemplatesResult(this.resourceTemplates, null));
}; };

View file

@ -23,6 +23,7 @@ import java.io.InputStreamReader;
import java.io.OutputStream; import java.io.OutputStream;
import java.nio.charset.StandardCharsets; import java.nio.charset.StandardCharsets;
import java.time.Duration; import java.time.Duration;
import java.util.concurrent.Executors;
import java.util.function.Function; import java.util.function.Function;
import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.core.type.TypeReference;
@ -97,8 +98,8 @@ public class StdioServerTransport implements McpTransport {
this.outputStream = System.out; this.outputStream = System.out;
// Use bounded schedulers for better resource management // Use bounded schedulers for better resource management
this.inboundScheduler = Schedulers.newBoundedElastic(1, 1, "inbound"); this.inboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), "inbound");
this.outboundScheduler = Schedulers.newBoundedElastic(1, 1, "outbound"); this.outboundScheduler = Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(), "outbound");
} }
@Override @Override
@ -139,38 +140,43 @@ public class StdioServerTransport implements McpTransport {
private void startInboundProcessing() { private void startInboundProcessing() {
this.inboundScheduler.schedule(() -> { this.inboundScheduler.schedule(() -> {
inboundReady.tryEmitValue(null); inboundReady.tryEmitValue(null);
try (BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) { BufferedReader reader = null;
String line; try {
while (!isClosing && (line = reader.readLine()) != null) { reader = new BufferedReader(new InputStreamReader(inputStream));
while (!isClosing) {
try { try {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line); String line = reader.readLine();
if (!this.inboundSink.tryEmitNext(message).isSuccess()) { if (line == null || isClosing) {
break;
}
try {
JSONRPCMessage message = McpSchema.deserializeJsonRpcMessage(this.objectMapper, line);
if (!this.inboundSink.tryEmitNext(message).isSuccess()) {
if (!isClosing) {
logger.error("Failed to enqueue message");
}
break;
}
}
catch (Exception e) {
if (!isClosing) { if (!isClosing) {
logger.error("Failed to enqueue message"); logger.error("Error processing inbound message", e);
} }
break; break;
} }
} }
catch (Exception e) { catch (IOException e) {
if (!isClosing) { if (!isClosing) {
logger.error("Error processing inbound message", e); logger.error("Error reading from stdin", e);
} }
break; break;
} }
} }
} }
catch (IOException e) { catch (Exception e) {
// Check isClosing before the error occurs to properly categorize it if (!isClosing) {
boolean wasClosing = isClosing; logger.error("Error in inbound processing", e);
isClosing = true;
if (!wasClosing && e.getMessage().equals("Pipe closed")) {
logger.debug("Stream closed during shutdown", e);
}
else if (!wasClosing) {
logger.error("Error reading from stdin", e);
}
else {
logger.debug("Stream error during shutdown", e);
} }
} }
finally { finally {
@ -234,6 +240,7 @@ public class StdioServerTransport implements McpTransport {
@Override @Override
public Mono<Void> closeGracefully() { public Mono<Void> closeGracefully() {
return Mono.fromRunnable(() -> { return Mono.fromRunnable(() -> {
isClosing = true; isClosing = true;
logger.debug("Initiating graceful shutdown"); logger.debug("Initiating graceful shutdown");
@ -244,18 +251,10 @@ public class StdioServerTransport implements McpTransport {
return Mono.delay(Duration.ofMillis(100)); return Mono.delay(Duration.ofMillis(100));
})).then(Mono.fromRunnable(() -> { })).then(Mono.fromRunnable(() -> {
try { try {
// Dispose schedulers first // Dispose schedulers with longer timeout
inboundScheduler.dispose(); inboundScheduler.dispose();
outboundScheduler.dispose(); outboundScheduler.dispose();
// Wait for schedulers to terminate
if (!inboundScheduler.isDisposed()) {
inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
}
if (!outboundScheduler.isDisposed()) {
outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
}
logger.info("Graceful shutdown completed"); logger.info("Graceful shutdown completed");
} }
catch (Exception e) { catch (Exception e) {

View file

@ -66,6 +66,7 @@ public abstract class AbstractMcpSyncServerTests {
@BeforeEach @BeforeEach
void setUp() { void setUp() {
// onStart();
} }
@AfterEach @AfterEach

View file

@ -17,6 +17,8 @@
package org.springframework.ai.mcp.server; package org.springframework.ai.mcp.server;
import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Timeout;
import org.springframework.http.server.reactive.HttpHandler; import org.springframework.http.server.reactive.HttpHandler;
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
import org.springframework.web.reactive.function.server.RouterFunctions; import org.springframework.web.reactive.function.server.RouterFunctions;
@ -31,6 +33,7 @@ import org.springframework.ai.mcp.spec.McpTransport;
* *
* @author Christian Tzolov * @author Christian Tzolov
*/ */
@Timeout(15) // Giving extra time beyond the client timeout
class SseMcpAsyncServerTests extends AbstractMcpAsyncServerTests { class SseMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
private static final int PORT = 8181; private static final int PORT = 8181;

View file

@ -17,6 +17,8 @@
package org.springframework.ai.mcp.server; package org.springframework.ai.mcp.server;
import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Timeout;
import org.springframework.http.server.reactive.HttpHandler; import org.springframework.http.server.reactive.HttpHandler;
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter; import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
import org.springframework.web.reactive.function.server.RouterFunctions; import org.springframework.web.reactive.function.server.RouterFunctions;
@ -31,6 +33,7 @@ import org.springframework.ai.mcp.spec.McpTransport;
* *
* @author Christian Tzolov * @author Christian Tzolov
*/ */
@Timeout(15) // Giving extra time beyond the client timeout
class SseMcpSyncServerTests extends AbstractMcpSyncServerTests { class SseMcpSyncServerTests extends AbstractMcpSyncServerTests {
private static final int PORT = 8182; private static final int PORT = 8182;

View file

@ -16,6 +16,8 @@
package org.springframework.ai.mcp.server; package org.springframework.ai.mcp.server;
import org.junit.jupiter.api.Timeout;
import org.springframework.ai.mcp.server.transport.StdioServerTransport; import org.springframework.ai.mcp.server.transport.StdioServerTransport;
import org.springframework.ai.mcp.spec.McpTransport; import org.springframework.ai.mcp.spec.McpTransport;
@ -24,6 +26,7 @@ import org.springframework.ai.mcp.spec.McpTransport;
* *
* @author Christian Tzolov * @author Christian Tzolov
*/ */
@Timeout(15) // Giving extra time beyond the client timeout
class StdioMcpAsyncServerTests extends AbstractMcpAsyncServerTests { class StdioMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
@Override @Override

View file

@ -16,6 +16,8 @@
package org.springframework.ai.mcp.server; package org.springframework.ai.mcp.server;
import org.junit.jupiter.api.Timeout;
import org.springframework.ai.mcp.server.transport.StdioServerTransport; import org.springframework.ai.mcp.server.transport.StdioServerTransport;
import org.springframework.ai.mcp.spec.McpTransport; import org.springframework.ai.mcp.spec.McpTransport;
@ -24,6 +26,7 @@ import org.springframework.ai.mcp.spec.McpTransport;
* *
* @author Christian Tzolov * @author Christian Tzolov
*/ */
@Timeout(15) // Giving extra time beyond the client timeout
class StdioMcpSyncServerTests extends AbstractMcpSyncServerTests { class StdioMcpSyncServerTests extends AbstractMcpSyncServerTests {
@Override @Override

View file

@ -16,7 +16,6 @@
package org.springframework.ai.mcp.server.transport; package org.springframework.ai.mcp.server.transport;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream; import java.io.ByteArrayOutputStream;
import java.io.InputStream; import java.io.InputStream;
import java.io.PrintStream; import java.io.PrintStream;
@ -51,8 +50,6 @@ class StdioServerTransportTests {
private ByteArrayOutputStream testOut; private ByteArrayOutputStream testOut;
private ByteArrayInputStream testIn;
private ByteArrayOutputStream testErr; private ByteArrayOutputStream testErr;
private PrintStream testOutPrintStream; private PrintStream testOutPrintStream;
@ -89,7 +86,6 @@ class StdioServerTransportTests {
void shouldHandleIncomingMessages() throws Exception { void shouldHandleIncomingMessages() throws Exception {
// Prepare test input // Prepare test input
String jsonMessage = "{\"jsonrpc\":\"2.0\",\"method\":\"test\",\"params\":{},\"id\":1}"; String jsonMessage = "{\"jsonrpc\":\"2.0\",\"method\":\"test\",\"params\":{},\"id\":1}";
testIn = new ByteArrayInputStream((jsonMessage + "\n").getBytes(StandardCharsets.UTF_8));
// Create transport with test streams // Create transport with test streams
transport = new StdioServerTransport(objectMapper); transport = new StdioServerTransport(objectMapper);

View file

@ -174,6 +174,7 @@
<!-- Output test execution times in the logs --> <!-- Output test execution times in the logs -->
<redirectTestOutputToFile>false</redirectTestOutputToFile> <redirectTestOutputToFile>false</redirectTestOutputToFile>
</configuration> </configuration>
</plugin> </plugin>
<plugin> <plugin>