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: simplify StdioServerTransport to only use System streams

Browse source 
Removes custom stream injection capability from StdioServerTransport and
simplifies it to only use System.in/out.
Updates related tests to use System streams instead of piped streams.
This commit is contained in:
Christian Tzolov 2024-12-30 11:52:43 +01:00
commit fcf0ff0ae2
4 changed files with 34 additions and 115 deletions
Download patch file Download diff file Expand all files Collapse all files

80
mcp/src/main/java/org/springframework/ai/mcp/server/transport/StdioServerTransport.java
View file

@ -77,7 +77,7 @@ public class StdioServerTransport implements McpTransport {
* Creates a new StdioServerTransport with a default ObjectMapper and System streams. * Creates a new StdioServerTransport with a default ObjectMapper and System streams.
*/ */
public StdioServerTransport() { public StdioServerTransport() {
this(new ObjectMapper(), System.in, System.out); this(new ObjectMapper());
} }
/** /**
@ -86,26 +86,15 @@ public class StdioServerTransport implements McpTransport {
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization * @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
*/ */
public StdioServerTransport(ObjectMapper objectMapper) { public StdioServerTransport(ObjectMapper objectMapper) {
this(objectMapper, System.in, System.out);
}
/**
* Creates a new StdioServerTransport with the specified ObjectMapper and streams.
* @param objectMapper The ObjectMapper to use for JSON serialization/deserialization
* @param inputStream The input stream to read from
* @param outputStream The output stream to write to
*/
public StdioServerTransport(ObjectMapper objectMapper, InputStream inputStream, OutputStream outputStream) {
Assert.notNull(objectMapper, "The ObjectMapper can not be null"); Assert.notNull(objectMapper, "The ObjectMapper can not be null");
Assert.notNull(inputStream, "The InputStream can not be null");
Assert.notNull(outputStream, "The OutputStream can not be null");
this.inboundSink = Sinks.many().unicast().onBackpressureBuffer(); this.inboundSink = Sinks.many().unicast().onBackpressureBuffer();
this.outboundSink = Sinks.many().unicast().onBackpressureBuffer(); this.outboundSink = Sinks.many().unicast().onBackpressureBuffer();
this.objectMapper = objectMapper; this.objectMapper = objectMapper;
this.inputStream = inputStream; this.inputStream = System.in;
this.outputStream = outputStream; 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.newBoundedElastic(1, 1, "inbound");
@ -248,50 +237,31 @@ public class StdioServerTransport implements McpTransport {
return Mono.fromRunnable(() -> { return Mono.fromRunnable(() -> {
isClosing = true; isClosing = true;
logger.debug("Initiating graceful shutdown"); logger.debug("Initiating graceful shutdown");
}) }).then(Mono.defer(() -> {
// .then(Mono.defer(() -> { // First complete the sinks to stop processing
// // First complete the sinks to stop processing inboundSink.tryEmitComplete();
// inboundSink.tryEmitComplete(); outboundSink.tryEmitComplete();
// outboundSink.tryEmitComplete(); return Mono.delay(Duration.ofMillis(100));
// return Mono.delay(Duration.ofMillis(100)); })).then(Mono.fromRunnable(() -> {
// })) try {
.then(Mono.fromRunnable(() -> { // Dispose schedulers first
try { inboundScheduler.dispose();
// Dispose schedulers first outboundScheduler.dispose();
inboundScheduler.dispose();
outboundScheduler.dispose();
// Wait for schedulers to terminate // Wait for schedulers to terminate
if (!inboundScheduler.isDisposed()) { if (!inboundScheduler.isDisposed()) {
inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5)); inboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
}
if (!outboundScheduler.isDisposed()) {
outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
}
// Only after schedulers are disposed, close the streams
try {
if (inputStream != System.in) {
inputStream.close();
}
if (outputStream != System.out) {
outputStream.flush();
outputStream.close();
}
}
catch (IOException e) {
// Log but don't throw since we're shutting down
logger.debug("Error closing streams during shutdown", e);
}
logger.info("Graceful shutdown completed");
} }
catch (Exception e) { if (!outboundScheduler.isDisposed()) {
logger.error("Error during graceful shutdown", e); outboundScheduler.disposeGracefully().block(Duration.ofSeconds(5));
} }
}))
.then() logger.info("Graceful shutdown completed");
.subscribeOn(Schedulers.boundedElastic()); }
catch (Exception e) {
logger.error("Error during graceful shutdown", e);
}
})).then().subscribeOn(Schedulers.boundedElastic());
} }
@Override @Override

4
mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpAsyncServerTests.java
View file

@ -16,8 +16,6 @@
package org.springframework.ai.mcp.server; package org.springframework.ai.mcp.server;
import com.fasterxml.jackson.databind.ObjectMapper;
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;
@ -30,7 +28,7 @@ class StdioMcpAsyncServerTests extends AbstractMcpAsyncServerTests {
@Override @Override
protected McpTransport createMcpTransport() { protected McpTransport createMcpTransport() {
return new StdioServerTransport(new ObjectMapper()); return new StdioServerTransport();
} }
} }

51
mcp/src/test/java/org/springframework/ai/mcp/server/StdioMcpSyncServerTests.java
View file

@ -16,11 +16,6 @@
package org.springframework.ai.mcp.server; package org.springframework.ai.mcp.server;
import java.io.PipedInputStream;
import java.io.PipedOutputStream;
import com.fasterxml.jackson.databind.ObjectMapper;
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;
@ -31,53 +26,9 @@ import org.springframework.ai.mcp.spec.McpTransport;
*/ */
class StdioMcpSyncServerTests extends AbstractMcpSyncServerTests { class StdioMcpSyncServerTests extends AbstractMcpSyncServerTests {
private PipedInputStream testInput;
private PipedOutputStream testOutput;
private PipedOutputStream writeToInput;
private PipedInputStream readFromOutput;
@Override @Override
protected McpTransport createMcpTransport() { protected McpTransport createMcpTransport() {
try { return new StdioServerTransport();
testInput = new PipedInputStream();
writeToInput = new PipedOutputStream(testInput);
testOutput = new PipedOutputStream();
readFromOutput = new PipedInputStream(testOutput);
return new StdioServerTransport(new ObjectMapper(), testInput, testOutput);
}
catch (Exception e) {
throw new RuntimeException("Failed to initialize test streams", e);
}
}
@Override
protected void onStart() {
// No special setup needed for stdio transport
}
@Override
protected void onClose() {
try {
if (testInput != null) {
testInput.close();
}
if (testOutput != null) {
testOutput.close();
}
if (writeToInput != null) {
writeToInput.close();
}
if (readFromOutput != null) {
readFromOutput.close();
}
}
catch (Exception e) {
// Log but don't throw since this is cleanup
e.printStackTrace();
}
} }
} }

12
mcp/src/test/java/org/springframework/ai/mcp/server/transport/StdioServerTransportTests.java
View file

@ -92,7 +92,7 @@ class StdioServerTransportTests {
testIn = new ByteArrayInputStream((jsonMessage + "\n").getBytes(StandardCharsets.UTF_8)); testIn = new ByteArrayInputStream((jsonMessage + "\n").getBytes(StandardCharsets.UTF_8));
// Create transport with test streams // Create transport with test streams
transport = new StdioServerTransport(objectMapper, testIn, testOutPrintStream); transport = new StdioServerTransport(objectMapper);
// Parse expected message // Parse expected message
McpSchema.JSONRPCRequest expected = objectMapper.readValue(jsonMessage, McpSchema.JSONRPCRequest.class); McpSchema.JSONRPCRequest expected = objectMapper.readValue(jsonMessage, McpSchema.JSONRPCRequest.class);
@ -109,9 +109,9 @@ class StdioServerTransportTests {
@Disabled @Disabled
void shouldHandleOutgoingMessages() throws Exception { void shouldHandleOutgoingMessages() throws Exception {
// Create transport with test streams // Create transport with test streams
transport = new StdioServerTransport(objectMapper, new BlockingInputStream(), testOutPrintStream); transport = new StdioServerTransport(objectMapper);
// transport = new StdioServerTransport(objectMapper, new ByteArrayInputStream(new // transport = new StdioServerTransport(objectMapper, new BlockingInputStream(),
// byte[0]), testOutPrintStream); // testOutPrintStream);
// Create test messages // Create test messages
JSONRPCRequest initMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "init", "init-id", JSONRPCRequest initMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "init", "init-id",
@ -137,7 +137,7 @@ class StdioServerTransportTests {
@Test @Test
void shouldWaitForProcessorsBeforeSendingMessage() { void shouldWaitForProcessorsBeforeSendingMessage() {
// Create transport with test streams // Create transport with test streams
transport = new StdioServerTransport(objectMapper, new ByteArrayInputStream(new byte[0]), testOutPrintStream); transport = new StdioServerTransport(objectMapper);
// Create test message // Create test message
JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test", "test-id", JSONRPCRequest testMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "test", "test-id",
@ -154,7 +154,7 @@ class StdioServerTransportTests {
@Test @Test
void shouldCloseGracefully() { void shouldCloseGracefully() {
// Create transport with test streams // Create transport with test streams
transport = new StdioServerTransport(objectMapper, new ByteArrayInputStream(new byte[0]), testOutPrintStream); transport = new StdioServerTransport(objectMapper);
// Create test message // Create test message
JSONRPCRequest initMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "init", "init-id", JSONRPCRequest initMessage = new JSONRPCRequest(McpSchema.JSONRPC_VERSION, "init", "init-id",