modelcontextprotocol/java-sdk
Watch
1
0
Fork
You've already forked java-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

non-blocking sinks processing on dedicated Schedulers

Browse source
This commit is contained in:
Dariusz Jędrzejczyk 2024-12-07 00:35:32 +01:00
commit 076c214400
No known key found for this signature in database
GPG key ID: FD786A4E2335C448
6 changed files with 122 additions and 73 deletions
Download patch file Download diff file Expand all files Collapse all files

8
src/main/java/spring/ai/mcp/client/McpAsyncClient.java
View file

@ -1,22 +1,20 @@
package spring.ai.mcp.client;
import java.time.Duration;
import java.util.concurrent.CompletableFuture;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.core.publisher.Mono;
import spring.ai.mcp.spec.McpAsyncTransport;
import spring.ai.mcp.spec.McpSchema;
import spring.ai.mcp.spec.McpSession;
import spring.ai.mcp.spec.McpTransport;
public class McpAsyncClient extends McpAsyncSession {
public McpAsyncClient(McpTransport transport) {
public McpAsyncClient(McpAsyncTransport transport) {
this(transport, Duration.ofSeconds(10), new ObjectMapper());
}
public McpAsyncClient(McpTransport transport, Duration requestTimeout, ObjectMapper objectMapper) {
public McpAsyncClient(McpAsyncTransport transport, Duration requestTimeout, ObjectMapper objectMapper) {
super(requestTimeout, objectMapper, transport);
}

22
src/main/java/spring/ai/mcp/client/McpAsyncSession.java
View file

@ -3,18 +3,15 @@ package spring.ai.mcp.client;
import java.time.Duration;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import reactor.core.publisher.Mono;
import reactor.core.publisher.MonoSink;
import reactor.core.publisher.Sinks;
import spring.ai.mcp.spec.McpSchema;
import spring.ai.mcp.spec.McpSession;
import spring.ai.mcp.spec.McpAsyncTransport;
import spring.ai.mcp.spec.McpTransport;
import spring.ai.mcp.spec.McpSchema;
public class McpAsyncSession {
@ -25,11 +22,11 @@ public class McpAsyncSession {
private final ObjectMapper objectMapper;
private final McpTransport transport;
private final McpAsyncTransport transport;
public McpAsyncSession(Duration requestTimeout,
ObjectMapper objectMapper,
McpTransport transport) {
McpAsyncTransport transport) {
this.requestTimeout = requestTimeout;
this.objectMapper = objectMapper;
this.transport = transport;
@ -65,8 +62,15 @@ public class McpAsyncSession {
this.pendingResponses.put(requestId, sink);
McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method, requestId, requestParams);
try {
// TODO convert the transport to non-blocking as well
this.transport.sendMessage(jsonrpcRequest);
// TODO: This is non-blocking, but it's actually a synchronous call,
// perhaps there's no need to make it return Mono?
this.transport.sendMessage(jsonrpcRequest)
// TODO: It's most efficient to create a dedicated
// Subscriber here
.subscribe(v -> {}, e -> {
this.pendingResponses.remove(requestId);
sink.error(e);
});
} catch (Exception e) {
sink.error(e);
}

10
src/main/java/spring/ai/mcp/client/McpClient.java
View file

@ -20,6 +20,8 @@ import java.util.concurrent.CompletableFuture;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import spring.ai.mcp.spec.McpAsyncTransport;
import spring.ai.mcp.spec.McpTransport;
import spring.ai.mcp.spec.McpSchema;
import spring.ai.mcp.spec.McpSession;
import spring.ai.mcp.spec.McpTransport;
@ -59,20 +61,20 @@ public class McpClient {
// JSONtoPOJOCodec type
// .sync();
public static McpAsyncClient async(McpTransport transport) {
public static McpAsyncClient async(McpAsyncTransport transport) {
return new McpAsyncClient(transport);
}
public static McpAsyncClient async(McpTransport transport, Duration requestTimeout,
public static McpAsyncClient async(McpAsyncTransport transport, Duration requestTimeout,
ObjectMapper objectMapper) {
return new McpAsyncClient(transport, requestTimeout, objectMapper);
}
public static McpSyncClient sync(McpTransport transport) {
public static McpSyncClient sync(McpAsyncTransport transport) {
return new McpSyncClient(async(transport));
}
public static McpSyncClient sync(McpTransport transport, Duration requestTimeout,
public static McpSyncClient sync(McpAsyncTransport transport, Duration requestTimeout,
ObjectMapper objectMapper) {
return new McpSyncClient(async(transport, requestTimeout, objectMapper));
}

82
src/main/java/spring/ai/mcp/client/stdio/StdioServerTransport.java
View file

@ -6,9 +6,10 @@ import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;
import spring.ai.mcp.spec.DefaultMcpTransport;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
import spring.ai.mcp.spec.McpSchema.JSONRPCNotification;
@ -31,7 +32,9 @@ public class StdioServerTransport extends DefaultMcpTransport {
private BufferedWriter processWriter;
private ExecutorService executorService;
private Scheduler inboundScheduler;
private Scheduler outboundScheduler;
private Scheduler errorScheduler;
private volatile boolean isRunning;
@ -49,7 +52,15 @@ public class StdioServerTransport extends DefaultMcpTransport {
this.params = params;
// Start threads
this.executorService = Executors.newFixedThreadPool(3);
this.inboundScheduler =
Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
"inbound");
this.outboundScheduler =
Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
"outbound");
this.errorScheduler =
Schedulers.fromExecutorService(Executors.newSingleThreadExecutor(),
"error");
}
@Override
@ -83,9 +94,9 @@ public class StdioServerTransport extends DefaultMcpTransport {
// Start threads
this.isRunning = true;
startReaderThread();
startWriterThread();
startErrorReaderThread();
startInboundProcessing();
startOutboundProcessing();
startErrorProcessing();
}
@ -98,9 +109,8 @@ public class StdioServerTransport extends DefaultMcpTransport {
}
}
private void startErrorReaderThread() {
this.executorService.execute(() -> {
private void startErrorProcessing() {
this.errorScheduler.schedule(() -> {
try {
String line;
while (isRunning && processErrorReader != null && (line = processErrorReader.readLine()) != null) {
@ -143,15 +153,15 @@ public class StdioServerTransport extends DefaultMcpTransport {
throw new IllegalArgumentException("Cannot deserialize JSONRPCMessage: " + jsonText);
}
private void startReaderThread() {
this.executorService.execute(() -> {
private void startInboundProcessing() {
this.inboundScheduler.schedule(() -> {
try {
String line;
while (this.isRunning && this.processReader != null && (line = this.processReader.readLine()) != null) {
try {
JSONRPCMessage message = deserializeJsonRpcMessage(line);
if (!this.getInboundSink().tryEmitNext(message).isSuccess()) {
// TODO: Back off, reschedule -> same comment as for writing
// TODO: Back off, reschedule, give up?
throw new RuntimeException("Failed to enqueue message");
}
}
@ -171,29 +181,33 @@ public class StdioServerTransport extends DefaultMcpTransport {
});
}
private void startWriterThread() {
this.executorService.execute(() -> {
this.getOutboundSink().asFlux().handle((message, s) -> {
if (message != null) {
try {
this.processWriter.write(ModelOptionsUtils.toJsonString(message));
this.processWriter.newLine();
this.processWriter.flush();
s.next(message);
} catch (IOException e) {
s.error(new RuntimeException(e));
}
}
}).subscribe();
});
private void startOutboundProcessing() {
this.getOutboundSink()
.asFlux()
// this bit is important since writes come from user threads and we
// want to ensure that the actual writing happens on a dedicated thread
.publishOn(outboundScheduler)
.handle((message, s) -> {
if (message != null) {
try {
this.processWriter.write(ModelOptionsUtils.toJsonString(message));
this.processWriter.newLine();
this.processWriter.flush();
s.next(message);
}
catch (IOException e) {
s.error(new RuntimeException(e));
}
}
})
.subscribe();
}
// TODO: provide a non-blocking variant with graceful option
public void stop() {
this.isRunning = false;
// Interrupt threads
this.executorService.shutdownNow();
this.outboundScheduler.dispose();
this.inboundScheduler.dispose();
this.errorScheduler.dispose();
// Destroy process
if (this.process != null) {
@ -201,10 +215,6 @@ public class StdioServerTransport extends DefaultMcpTransport {
}
}
public boolean isRunning() {
return this.isRunning;
}
@Override
public void close() {
stop();

34
src/main/java/spring/ai/mcp/spec/DefaultMcpTransport.java
View file

@ -16,10 +16,9 @@
package spring.ai.mcp.spec;
import java.time.Duration;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.function.Consumer;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
@ -27,7 +26,7 @@ import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
* @author Christian Tzolov
* @since 1.0.0
*/
public class DefaultMcpTransport implements McpTransport {
public class DefaultMcpTransport implements McpAsyncTransport {
private final Sinks.Many<String> errorSink;
@ -35,8 +34,6 @@ public class DefaultMcpTransport implements McpTransport {
private final Sinks.Many<JSONRPCMessage> outboundSink;
private final ExecutorService executorService;
private final Duration writeTimeout;
private Consumer<JSONRPCMessage> messageReader = message -> System.out.println("Received message: " + message);
@ -56,23 +53,18 @@ public class DefaultMcpTransport implements McpTransport {
this.outboundSink = Sinks.many().unicast().onBackpressureBuffer();
this.writeTimeout = readTimeout;
this.executorService = Executors.newFixedThreadPool(2);
this.handleIncomingMessages();
}
private void handleIncomingMessages() {
this.executorService.execute(() -> {
this.inboundSink.asFlux()
.subscribe(message -> this.messageReader.accept(message));
});
this.inboundSink.asFlux()
.subscribe(message -> this.messageReader.accept(message));
}
private void handleIncomingErrors() {
this.executorService.execute(() -> {
this.errorSink.asFlux().subscribe(e -> {
this.errorReader.accept(e);
System.err.println(e);
});
this.errorSink.asFlux().subscribe(e -> {
this.errorReader.accept(e);
System.err.println(e);
});
}
@ -106,15 +98,19 @@ public class DefaultMcpTransport implements McpTransport {
@Override
public void close() {
// this.onClose();
this.executorService.shutdownNow();
}
public void sendMessage(JSONRPCMessage message) {
if (!this.outboundSink.tryEmitNext(message).isSuccess()) {
@Override
public Mono<Void> sendMessage(JSONRPCMessage message) {
if (this.outboundSink.tryEmitNext(message).isSuccess()) {
// TODO: essentially we could reschedule ourselves in some time and make
// another attempt with the already read data but pause reading until
// success
throw new RuntimeException("Failed to enqueue message");
// In this approach we delegate the retry and the backpressure onto the
// caller. This might be enough for most cases.
return Mono.empty();
} else {
return Mono.error(new RuntimeException("Failed to enqueue message"));
}
}
}

39
src/main/java/spring/ai/mcp/spec/McpAsyncTransport.java Normal file
View file

@ -0,0 +1,39 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package spring.ai.mcp.spec;
import java.util.function.Consumer;
import reactor.core.publisher.Mono;
import spring.ai.mcp.spec.McpSchema.JSONRPCMessage;
/**
* @author Christian Tzolov
* @since 1.0.0
*/
public interface McpAsyncTransport extends AutoCloseable{
void start();
void close();
void setMessageHandler(Consumer<JSONRPCMessage> messageReader);
void setErrorHandler(Consumer<String> errorReader);
Mono<Void> sendMessage(JSONRPCMessage message);
}