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:
parent
fcf0ff0ae2
commit
ad3cdd20a9
9 changed files with 50 additions and 41 deletions
|
|
@ -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));
|
||||||
};
|
};
|
||||||
|
|
|
||||||
|
|
@ -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) {
|
||||||
|
|
|
||||||
|
|
@ -66,6 +66,7 @@ public abstract class AbstractMcpSyncServerTests {
|
||||||
|
|
||||||
@BeforeEach
|
@BeforeEach
|
||||||
void setUp() {
|
void setUp() {
|
||||||
|
// onStart();
|
||||||
}
|
}
|
||||||
|
|
||||||
@AfterEach
|
@AfterEach
|
||||||
|
|
|
||||||
|
|
@ -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;
|
||||||
|
|
|
||||||
|
|
@ -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;
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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);
|
||||||
|
|
|
||||||
1
pom.xml
1
pom.xml
|
|
@ -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>
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue