Bound STDIO server reads
Signed-off-by: Daniel Garnier-Moiroux <git@garnier.wf>
This commit is contained in:
parent
ef037dae7b
commit
3b72eec0b8
1 changed files with 101 additions and 3 deletions
|
|
@ -41,6 +41,8 @@ import reactor.core.scheduler.Schedulers;
|
|||
*/
|
||||
public class StdioServerTransportProvider implements McpServerTransportProvider {
|
||||
|
||||
private static final int DEFAULT_INPUT_MAX_SIZE = 16 * 1024 * 1024; // 16MB
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(StdioServerTransportProvider.class);
|
||||
|
||||
private final McpJsonMapper jsonMapper;
|
||||
|
|
@ -49,6 +51,8 @@ public class StdioServerTransportProvider implements McpServerTransportProvider
|
|||
|
||||
private final OutputStream outputStream;
|
||||
|
||||
private final int inputMaxSize;
|
||||
|
||||
private McpServerSession session;
|
||||
|
||||
private final AtomicBoolean isClosing = new AtomicBoolean(false);
|
||||
|
|
@ -72,13 +76,29 @@ public class StdioServerTransportProvider implements McpServerTransportProvider
|
|||
* @param outputStream The output stream to write to
|
||||
*/
|
||||
public StdioServerTransportProvider(McpJsonMapper jsonMapper, InputStream inputStream, OutputStream outputStream) {
|
||||
this(jsonMapper, inputStream, outputStream, DEFAULT_INPUT_MAX_SIZE);
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a new StdioServerTransportProvider.
|
||||
* @param jsonMapper The JsonMapper to use for JSON serialization/deserialization
|
||||
* @param inputStream The input stream to read from
|
||||
* @param outputStream The output stream to write to
|
||||
* @param inputMaxSize The maximum number of characters read for a single inbound
|
||||
* message. A peer that sends a longer message (or never terminates a line) has its
|
||||
* message rejected instead of forcing the transport to buffer it in memory.
|
||||
*/
|
||||
public StdioServerTransportProvider(McpJsonMapper jsonMapper, InputStream inputStream, OutputStream outputStream,
|
||||
int inputMaxSize) {
|
||||
Assert.notNull(jsonMapper, "The JsonMapper can not be null");
|
||||
Assert.notNull(inputStream, "The InputStream can not be null");
|
||||
Assert.notNull(outputStream, "The OutputStream can not be null");
|
||||
Assert.isTrue(inputMaxSize > 0, "inputMaxSize must be positive");
|
||||
|
||||
this.jsonMapper = jsonMapper;
|
||||
this.inputStream = inputStream;
|
||||
this.outputStream = outputStream;
|
||||
this.inputMaxSize = inputMaxSize;
|
||||
}
|
||||
|
||||
@Override
|
||||
|
|
@ -211,7 +231,7 @@ public class StdioServerTransportProvider implements McpServerTransportProvider
|
|||
reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
|
||||
while (!isClosing.get()) {
|
||||
try {
|
||||
String line = reader.readLine();
|
||||
String line = readLine(reader, inputMaxSize);
|
||||
if (line == null || isClosing.get()) {
|
||||
break;
|
||||
}
|
||||
|
|
@ -232,6 +252,10 @@ public class StdioServerTransportProvider implements McpServerTransportProvider
|
|||
break;
|
||||
}
|
||||
}
|
||||
catch (MaxSizeExceededException e) {
|
||||
logIfNotClosing("Inbound message exceeds the maximum allowed size", e);
|
||||
break;
|
||||
}
|
||||
catch (IOException e) {
|
||||
logIfNotClosing("Error reading from stdin", e);
|
||||
break;
|
||||
|
|
@ -304,6 +328,36 @@ public class StdioServerTransportProvider implements McpServerTransportProvider
|
|||
outboundConsumer.apply(outboundSink.asFlux()).subscribe();
|
||||
} // @formatter:on
|
||||
|
||||
/**
|
||||
* Read line with a max size.
|
||||
*/
|
||||
private static String readLine(BufferedReader reader, int maxSize)
|
||||
throws IOException, MaxSizeExceededException {
|
||||
StringBuilder sb = new StringBuilder();
|
||||
int c;
|
||||
while ((c = reader.read()) != -1) {
|
||||
if (c == '\n') {
|
||||
return sb.toString();
|
||||
}
|
||||
if (c == '\r') {
|
||||
// Consume an optional trailing '\n' so that "\r\n" is treated as a
|
||||
// single terminator, mirroring BufferedReader#readLine().
|
||||
reader.mark(1);
|
||||
int next = reader.read();
|
||||
if (next != '\n' && next != -1) {
|
||||
reader.reset();
|
||||
}
|
||||
return sb.toString();
|
||||
}
|
||||
if (sb.length() >= maxSize) {
|
||||
throw new MaxSizeExceededException(
|
||||
"Inbound message exceeds the maximum allowed size of " + maxSize + " characters");
|
||||
}
|
||||
sb.append((char) c);
|
||||
}
|
||||
return sb.isEmpty() ? null : sb.toString();
|
||||
}
|
||||
|
||||
private void logIfNotClosing(String message, Exception e) {
|
||||
if (!isClosing.get()) {
|
||||
logger.error(message, e);
|
||||
|
|
|
|||
Loading…
Reference in a new issue