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

feat(McpAsyncServer): Add non-blocking execution for tools and resources

Browse source 
- 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:
Christian Tzolov 2025-01-04 17:49:39 +01:00 • committed by Christian Tzolov
commit 580f95526f
2 changed files with 132 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

15
mcp/src/main/java/org/springframework/ai/mcp/server/McpAsyncServer.java
View file

@ -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()));

112
mcp/src/test/java/org/springframework/ai/mcp/server/SseAsyncIntegrationTests.java
View file

@ -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();
}
}