Fix stateless MCP tools calling (spring-projects#4215)

Signed-off-by: Daniel Garnier-Moiroux <[email protected]>
Signed-off-by: 家娃 <[email protected]>

Author: Kehrlann

auto-configurations/mcp/spring-ai-autoconfigure-mcp-server-webflux/src/test/java/org/springframework/ai/mcp/server/autoconfigure/StatelessWebClientWebFluxServerIT.java

@@ -25,6 +25,8 @@
 import java.util.Map;
 
 import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.model.ToolContext;
+import org.springframework.ai.mcp.McpToolUtils;
 import org.springframework.ai.mcp.client.common.autoconfigure.McpClientAutoConfiguration;
 import org.springframework.ai.mcp.client.common.autoconfigure.McpToolCallbackAutoConfiguration;
 import org.springframework.ai.mcp.client.webflux.autoconfigure.StreamableHttpWebFluxTransportAutoConfiguration;
@@ -33,6 +35,7 @@
 import org.springframework.ai.mcp.server.common.autoconfigure.properties.McpServerStreamableHttpProperties;
 import org.springframework.ai.mcp.server.common.autoconfigure.McpServerStatelessAutoConfiguration;
 import org.springframework.ai.mcp.server.common.autoconfigure.StatelessToolCallbackConverterAutoConfiguration;
+import org.springframework.ai.tool.function.FunctionToolCallback;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
@@ -123,7 +126,7 @@ void clientServerCapabilities() {
 						// TOOLS / SAMPLING / ELICITATION
 
 						// tool list
-						assertThat(mcpClient.listTools().tools()).hasSize(2);
+						assertThat(mcpClient.listTools().tools()).hasSize(3);
 						assertThat(mcpClient.listTools().tools())
 							.contains(Tool.builder().name("tool1").description("tool1 description").inputSchema("""
 									{
@@ -166,6 +169,18 @@ void clientServerCapabilities() {
 							.isEqualTo(json("""
 									{"result":5.0,"operation":"2 + 3","timestamp":"2024-01-01T10:00:00Z"}"""));
 
+						// TOOL FROM MCP TOOL UTILS
+						// Call the tool to ensure arguments are passed correctly
+						CallToolResult toUpperCaseResponse = mcpClient
+							.callTool(new McpSchema.CallToolRequest("toUpperCase", Map.of("input", "hello world")));
+						assertThat(toUpperCaseResponse).isNotNull();
+						assertThat(toUpperCaseResponse.isError()).isFalse();
+						assertThat(toUpperCaseResponse.content()).hasSize(1)
+							.first()
+							.isInstanceOf(TextContent.class)
+							.extracting("text")
+							.isEqualTo("\"HELLO WORLD\"");
+
 						// PROMPT / COMPLETION
 
 						// list prompts
@@ -254,7 +269,20 @@ public List<McpStatelessServerFeatures.SyncToolSpecification> myTools() {
 				})
 				.build();
 
-			return List.of(tool1, tool2);
+			// Tool 3
+
+			// Using a tool with McpToolUtils
+			McpStatelessServerFeatures.SyncToolSpecification tool3 = McpToolUtils
+				.toStatelessSyncToolSpecification(FunctionToolCallback
+					.builder("toUpperCase", (ToUpperCaseRequest req, ToolContext context) -> req.input().toUpperCase())
+					.description("Sets the input string to upper case")
+					.inputType(ToUpperCaseRequest.class)
+					.build(), null);
+
+			return List.of(tool1, tool2, tool3);
+		}
+
+		record ToUpperCaseRequest(String input) {
 		}
 
 		@Bean

mcp/common/src/main/java/org/springframework/ai/mcp/McpToolUtils.java

@@ -19,6 +19,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.function.BiFunction;
 
 import com.fasterxml.jackson.annotation.JsonAlias;
 import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
@@ -29,7 +30,9 @@
 import io.modelcontextprotocol.server.McpServerFeatures.AsyncToolSpecification;
 import io.modelcontextprotocol.server.McpStatelessServerFeatures;
 import io.modelcontextprotocol.server.McpSyncServerExchange;
+import io.modelcontextprotocol.server.McpTransportContext;
 import io.modelcontextprotocol.spec.McpSchema;
+import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
 import io.modelcontextprotocol.spec.McpSchema.Role;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Schedulers;
@@ -152,16 +155,6 @@ public static McpServerFeatures.SyncToolSpecification toSyncToolSpecification(To
 	 * Converts a Spring AI ToolCallback to an MCP SyncToolSpecification. This enables
 	 * Spring AI functions to be exposed as MCP tools that can be discovered and invoked
 	 * by language models.
-	 *
-	 * <p>
-	 * The conversion process:
-	 * <ul>
-	 * <li>Creates an MCP Tool with the function's name and input schema</li>
-	 * <li>Wraps the function's execution in a SyncToolSpecification that handles the MCP
-	 * protocol</li>
-	 * <li>Provides error handling and result formatting according to MCP
-	 * specifications</li>
-	 * </ul>
 	 * @param toolCallback the Spring AI function callback to convert
 	 * @param mimeType the MIME type of the output content
 	 * @return an MCP SyncToolSpecification that wraps the function callback
@@ -170,39 +163,50 @@ public static McpServerFeatures.SyncToolSpecification toSyncToolSpecification(To
 	public static McpServerFeatures.SyncToolSpecification toSyncToolSpecification(ToolCallback toolCallback,
 			MimeType mimeType) {
 
-		var tool = new McpSchema.Tool(toolCallback.getToolDefinition().name(),
-				toolCallback.getToolDefinition().description(), toolCallback.getToolDefinition().inputSchema());
+		SharedSyncToolSpecification sharedSpec = toSharedSyncToolSpecification(toolCallback, mimeType);
 
-		return new McpServerFeatures.SyncToolSpecification(tool, (exchange, request) -> {
-			try {
-				String callResult = toolCallback.call(ModelOptionsUtils.toJsonString(request),
-						new ToolContext(Map.of(TOOL_CONTEXT_MCP_EXCHANGE_KEY, exchange)));
-				if (mimeType != null && mimeType.toString().startsWith("image")) {
-					return new McpSchema.CallToolResult(List
-						.of(new McpSchema.ImageContent(List.of(Role.ASSISTANT), null, callResult, mimeType.toString())),
-							false);
-				}
-				return new McpSchema.CallToolResult(List.of(new McpSchema.TextContent(callResult)), false);
-			}
-			catch (Exception e) {
-				return new McpSchema.CallToolResult(List.of(new McpSchema.TextContent(e.getMessage())), true);
-			}
-		});
+		return new McpServerFeatures.SyncToolSpecification(sharedSpec.tool(), null,
+				(exchange, request) -> sharedSpec.sharedHandler().apply(exchange, request));
 	}
 
+	/**
+	 * Converts a Spring AI ToolCallback to an MCP StatelessSyncToolSpecification. This
+	 * enables Spring AI functions to be exposed as MCP tools that can be discovered and
+	 * invoked by language models.
+	 *
+	 * You can use the ToolCallback builder to create a new instance of ToolCallback using
+	 * either java.util.function.Function or Method reference.
+	 * @param toolCallback the Spring AI function callback to convert
+	 * @param mimeType the MIME type of the output content
+	 * @return an MCP StatelessSyncToolSpecification that wraps the function callback
+	 * @throws RuntimeException if there's an error during the function execution
+	 */
 	public static McpStatelessServerFeatures.SyncToolSpecification toStatelessSyncToolSpecification(
 			ToolCallback toolCallback, MimeType mimeType) {
 
+		var sharedSpec = toSharedSyncToolSpecification(toolCallback, mimeType);
+
+		return new McpStatelessServerFeatures.SyncToolSpecification(sharedSpec.tool(),
+				(exchange, request) -> sharedSpec.sharedHandler().apply(exchange, request));
+	}
+
+	private record SharedSyncToolSpecification(McpSchema.Tool tool,
+			BiFunction<Object, CallToolRequest, McpSchema.CallToolResult> sharedHandler) {
+	}
+
+	private static SharedSyncToolSpecification toSharedSyncToolSpecification(ToolCallback toolCallback,
+			MimeType mimeType) {
+
 		var tool = McpSchema.Tool.builder()
 			.name(toolCallback.getToolDefinition().name())
 			.description(toolCallback.getToolDefinition().description())
 			.inputSchema(toolCallback.getToolDefinition().inputSchema())
 			.build();
 
-		return new McpStatelessServerFeatures.SyncToolSpecification(tool, (mcpTransportContext, request) -> {
+		return new SharedSyncToolSpecification(tool, (exchangeOrContext, request) -> {
 			try {
-				String callResult = toolCallback.call(ModelOptionsUtils.toJsonString(request),
-						new ToolContext(Map.of(TOOL_CONTEXT_MCP_EXCHANGE_KEY, mcpTransportContext)));
+				String callResult = toolCallback.call(ModelOptionsUtils.toJsonString(request.arguments()),
+						new ToolContext(Map.of(TOOL_CONTEXT_MCP_EXCHANGE_KEY, exchangeOrContext)));
 				if (mimeType != null && mimeType.toString().startsWith("image")) {
 					return new McpSchema.CallToolResult(List
 						.of(new McpSchema.ImageContent(List.of(Role.ASSISTANT), null, callResult, mimeType.toString())),
@@ -329,8 +333,8 @@ public static McpStatelessServerFeatures.AsyncToolSpecification toStatelessAsync
 				toolCallback, mimeType);
 
 		return new McpStatelessServerFeatures.AsyncToolSpecification(statelessSyncToolSpecification.tool(),
-				(context, map) -> Mono
-					.fromCallable(() -> statelessSyncToolSpecification.callHandler().apply(context, map))
+				(context, request) -> Mono
+					.fromCallable(() -> statelessSyncToolSpecification.callHandler().apply(context.copy(), request))
 					.subscribeOn(Schedulers.boundedElastic()));
 	}