feat(McpAsyncServer): Add non-blocking execution for tools and resources
- Execute tool calls, resource reads and prompt handling in a non-blocking manner using Schedulers.boundedElastic(). This prevents blocking operations from impacting server responsiveness. - Added integration tests to verify non-blocking behavior with tools that make HTTP calls to external services. Related to #48 This is a temp patch until #48 is resolved properly.
This commit is contained in:
parent
b2082a0506
commit
580f95526f
8 changed files with 132 additions and 15 deletions
|
|
@ -32,14 +32,14 @@ Add the following dependencies to your Maven project:
|
|||
<dependency>
|
||||
<groupId>org.springframework.experimental</groupId>
|
||||
<artifactId>mcp</artifactId>
|
||||
<version>0.4.0-SNAPSHOT</version>
|
||||
<version>0.5.0-SNAPSHOT</version>
|
||||
</dependency>
|
||||
|
||||
<!-- For Spring AI integration -->
|
||||
<dependency>
|
||||
<groupId>org.springframework.experimental</groupId>
|
||||
<artifactId>spring-ai-mcp</artifactId>
|
||||
<version>0.4.0-SNAPSHOT</version>
|
||||
<version>0.5.0-SNAPSHOT</version>
|
||||
</dependency>
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ Add the following dependency to your Maven project:
|
|||
<dependency>
|
||||
<groupId>org.springframework.experimental</groupId>
|
||||
<artifactId>mcp</artifactId>
|
||||
<version>0.4.0-SNAPSHOT</version>
|
||||
<version>0.5.0-SNAPSHOT</version>
|
||||
</dependency>
|
||||
----
|
||||
|
||||
|
|
|
|||
|
|
@ -40,7 +40,7 @@ Maven::
|
|||
<dependency>
|
||||
<groupId>org.springframework.experimental</groupId>
|
||||
<artifactId>mcp</artifactId>
|
||||
<version>0.4.0-SNAPSHOT</version>
|
||||
<version>0.5.0-SNAPSHOT</version>
|
||||
</dependency>
|
||||
----
|
||||
+
|
||||
|
|
@ -50,7 +50,7 @@ Maven::
|
|||
<dependency>
|
||||
<groupId>org.springframework.experimental</groupId>
|
||||
<artifactId>spring-ai-mcp</artifactId>
|
||||
<version>0.4.0-SNAPSHOT</version>
|
||||
<version>0.5.0-SNAPSHOT</version>
|
||||
</dependency>
|
||||
----
|
||||
+
|
||||
|
|
|
|||
|
|
@ -36,6 +36,6 @@ To use this module, add the following dependency to your Maven project:
|
|||
<dependency>
|
||||
<groupId>org.springframework.experimental</groupId>
|
||||
<artifactId>spring-ai-mcp</artifactId>
|
||||
<version>0.4.0-SNAPSHOT</version>
|
||||
<version>0.5.0-SNAPSHOT</version>
|
||||
</dependency>
|
||||
----
|
||||
|
|
|
|||
|
|
@ -374,9 +374,9 @@ public class McpAsyncServer {
|
|||
return Mono.<Object>error(new McpError("Tool not found: " + callToolRequest.name()));
|
||||
}
|
||||
|
||||
CallToolResult callResponse = toolRegistration.get().call().apply(callToolRequest.arguments());
|
||||
|
||||
return Mono.just(callResponse);
|
||||
return Mono.fromCallable(() -> toolRegistration.get().call().apply(callToolRequest.arguments()))
|
||||
.map(result -> (Object) result)
|
||||
.subscribeOn(Schedulers.boundedElastic());
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -462,7 +462,9 @@ public class McpAsyncServer {
|
|||
});
|
||||
var resourceUri = resourceRequest.uri();
|
||||
if (this.resources.containsKey(resourceUri)) {
|
||||
return Mono.just(this.resources.get(resourceUri).readHandler().apply(resourceRequest));
|
||||
return Mono.fromCallable(() -> this.resources.get(resourceUri).readHandler().apply(resourceRequest))
|
||||
.map(result -> (Object) result)
|
||||
.subscribeOn(Schedulers.boundedElastic());
|
||||
}
|
||||
return Mono.error(new McpError("Resource not found: " + resourceUri));
|
||||
};
|
||||
|
|
@ -558,7 +560,10 @@ public class McpAsyncServer {
|
|||
|
||||
// Implement prompt retrieval logic here
|
||||
if (this.prompts.containsKey(promptRequest.name())) {
|
||||
return Mono.just(this.prompts.get(promptRequest.name()).promptHandler().apply(promptRequest));
|
||||
return Mono
|
||||
.fromCallable(() -> this.prompts.get(promptRequest.name()).promptHandler().apply(promptRequest))
|
||||
.map(result -> (Object) result)
|
||||
.subscribeOn(Schedulers.boundedElastic());
|
||||
}
|
||||
|
||||
return Mono.error(new McpError("Prompt not found: " + promptRequest.name()));
|
||||
|
|
|
|||
|
|
@ -31,17 +31,22 @@ import reactor.test.StepVerifier;
|
|||
|
||||
import org.springframework.ai.mcp.client.McpClient;
|
||||
import org.springframework.ai.mcp.client.transport.SseClientTransport;
|
||||
import org.springframework.ai.mcp.server.McpServer.ToolRegistration;
|
||||
import org.springframework.ai.mcp.server.transport.SseServerTransport;
|
||||
import org.springframework.ai.mcp.spec.McpError;
|
||||
import org.springframework.ai.mcp.spec.McpSchema;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CallToolResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ClientCapabilities;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CreateMessageRequest;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.CreateMessageResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.InitializeResult;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Role;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Root;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.ServerCapabilities;
|
||||
import org.springframework.ai.mcp.spec.McpSchema.Tool;
|
||||
import org.springframework.http.server.reactive.HttpHandler;
|
||||
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
|
||||
import org.springframework.web.client.RestClient;
|
||||
import org.springframework.web.reactive.function.client.WebClient;
|
||||
import org.springframework.web.reactive.function.server.RouterFunctions;
|
||||
|
||||
|
|
@ -314,4 +319,111 @@ public class SseAsyncIntegrationTests {
|
|||
mcpClient.close();
|
||||
}
|
||||
|
||||
// ---------------------------------------
|
||||
// Tools Tests
|
||||
// ---------------------------------------
|
||||
@Test
|
||||
void testToolCallSuccess() {
|
||||
|
||||
var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
|
||||
ToolRegistration tool1 = new ToolRegistration(
|
||||
new McpSchema.Tool("tool1", "tool1 description", Map.of("city", "String")), request -> {
|
||||
// perform a blocking call to a remote service
|
||||
String response = RestClient.create()
|
||||
.get()
|
||||
.uri("https://github.com/spring-projects-experimental/spring-ai-mcp/blob/main/README.md")
|
||||
.retrieve()
|
||||
.body(String.class);
|
||||
assertThat(response).isNotBlank();
|
||||
return callResponse;
|
||||
});
|
||||
|
||||
var mcpServer = McpServer.using(mcpServerTransport)
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.tools(tool1)
|
||||
.sync();
|
||||
|
||||
var mcpClient = clientBuilder.sync();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
assertThat(mcpClient.listTools().tools()).contains(tool1.tool());
|
||||
|
||||
CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of()));
|
||||
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response).isEqualTo(callResponse);
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testToolListChangeHandlingSuccess() {
|
||||
|
||||
var callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), null);
|
||||
ToolRegistration tool1 = new ToolRegistration(
|
||||
new McpSchema.Tool("tool1", "tool1 description", Map.of("city", "String")), request -> {
|
||||
// perform a blocking call to a remote service
|
||||
String response = RestClient.create()
|
||||
.get()
|
||||
.uri("https://github.com/spring-projects-experimental/spring-ai-mcp/blob/main/README.md")
|
||||
.retrieve()
|
||||
.body(String.class);
|
||||
assertThat(response).isNotBlank();
|
||||
return callResponse;
|
||||
});
|
||||
|
||||
var mcpServer = McpServer.using(mcpServerTransport)
|
||||
.capabilities(ServerCapabilities.builder().tools(true).build())
|
||||
.tools(tool1)
|
||||
.sync();
|
||||
|
||||
AtomicReference<List<Tool>> rootsRef = new AtomicReference<>();
|
||||
var mcpClient = clientBuilder.toolsChangeConsumer(toolsUpdate -> {
|
||||
// perform a blocking call to a remote service
|
||||
String response = RestClient.create()
|
||||
.get()
|
||||
.uri("https://github.com/spring-projects-experimental/spring-ai-mcp/blob/main/README.md")
|
||||
.retrieve()
|
||||
.body(String.class);
|
||||
assertThat(response).isNotBlank();
|
||||
rootsRef.set(toolsUpdate);
|
||||
}).sync();
|
||||
|
||||
InitializeResult initResult = mcpClient.initialize();
|
||||
assertThat(initResult).isNotNull();
|
||||
|
||||
assertThat(rootsRef.get()).isNull();
|
||||
|
||||
assertThat(mcpClient.listTools().tools()).contains(tool1.tool());
|
||||
|
||||
mcpServer.notifyToolsListChanged();
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).containsAll(List.of(tool1.tool()));
|
||||
});
|
||||
|
||||
// Remove a tool
|
||||
mcpServer.removeTool("tool1");
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).isEmpty();
|
||||
});
|
||||
|
||||
// Add a new tool
|
||||
ToolRegistration tool2 = new ToolRegistration(
|
||||
new McpSchema.Tool("tool2", "tool2 description", Map.of("city", "String")), request -> callResponse);
|
||||
|
||||
mcpServer.addTool(tool2);
|
||||
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(rootsRef.get()).containsAll(List.of(tool2.tool()));
|
||||
});
|
||||
|
||||
mcpClient.close();
|
||||
mcpServer.close();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ The server can be started in two transport modes, controlled by the `transport.m
|
|||
### Stdio Mode (Default)
|
||||
|
||||
```bash
|
||||
java -Dtransport.mode=stdio -jar target/spring-ai-mcp-sample-0.4.0-SNAPSHOT.jar
|
||||
java -Dtransport.mode=stdio -jar target/spring-ai-mcp-sample-0.5.0-SNAPSHOT.jar
|
||||
```
|
||||
|
||||
The Stdio mode server is automatically started by the client - no explicit server startup is needed.
|
||||
|
|
@ -38,7 +38,7 @@ In Stdio mode the server must not emit any messages/logs to the console (e.g. st
|
|||
|
||||
### SSE Mode
|
||||
```bash
|
||||
java -Dtransport.mode=sse -jar target/spring-ai-mcp-sample-0.4.0-SNAPSHOT.jar
|
||||
java -Dtransport.mode=sse -jar target/spring-ai-mcp-sample-0.5.0-SNAPSHOT.jar
|
||||
```
|
||||
|
||||
## Sample Clients
|
||||
|
|
@ -49,7 +49,7 @@ The project includes example clients for both transport modes:
|
|||
```java
|
||||
var stdioParams = ServerParameters.builder("java")
|
||||
.args("-Dtransport.mode=stdio", "-jar",
|
||||
"target/spring-ai-mcp-sample-0.4.0-SNAPSHOT.jar")
|
||||
"target/spring-ai-mcp-sample-0.5.0-SNAPSHOT.jar")
|
||||
.build();
|
||||
|
||||
var transport = new StdioClientTransport(stdioParams);
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ public class ClientStdio {
|
|||
|
||||
var stdioParams = ServerParameters.builder("java")
|
||||
.args("-Dtransport.mode=stdio", "-jar",
|
||||
"spring-ai-mcp-sample/target/spring-ai-mcp-sample-0.4.0-SNAPSHOT.jar")
|
||||
"spring-ai-mcp-sample/target/spring-ai-mcp-sample-0.5.0-SNAPSHOT.jar")
|
||||
.build();
|
||||
|
||||
var transport = new StdioClientTransport(stdioParams);
|
||||
|
|
|
|||
Loading…
Reference in a new issue