Skip to content

Commit 1df5f5c

Browse files
Kehrlanntzolov
Kehrlann
authored and
tzolov
committed
Fix stateless MCP tools calling (#4215)
Signed-off-by: Daniel Garnier-Moiroux <[email protected]>
1 parent e119eb7 commit 1df5f5c

File tree

2 files changed

+66
-34
lines changed

2 files changed

+66
-34
lines changed

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

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@
2525
import java.util.Map;
2626

2727
import org.junit.jupiter.api.Test;
28+
import org.springframework.ai.chat.model.ToolContext;
29+
import org.springframework.ai.mcp.McpToolUtils;
2830
import org.springframework.ai.mcp.client.common.autoconfigure.McpClientAutoConfiguration;
2931
import org.springframework.ai.mcp.client.common.autoconfigure.McpToolCallbackAutoConfiguration;
3032
import org.springframework.ai.mcp.client.webflux.autoconfigure.StreamableHttpWebFluxTransportAutoConfiguration;
@@ -33,6 +35,7 @@
3335
import org.springframework.ai.mcp.server.common.autoconfigure.properties.McpServerStreamableHttpProperties;
3436
import org.springframework.ai.mcp.server.common.autoconfigure.McpServerStatelessAutoConfiguration;
3537
import org.springframework.ai.mcp.server.common.autoconfigure.StatelessToolCallbackConverterAutoConfiguration;
38+
import org.springframework.ai.tool.function.FunctionToolCallback;
3639
import org.springframework.beans.factory.ObjectProvider;
3740
import org.springframework.boot.autoconfigure.AutoConfigurations;
3841
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
@@ -123,7 +126,7 @@ void clientServerCapabilities() {
123126
// TOOLS / SAMPLING / ELICITATION
124127

125128
// tool list
126-
assertThat(mcpClient.listTools().tools()).hasSize(2);
129+
assertThat(mcpClient.listTools().tools()).hasSize(3);
127130
assertThat(mcpClient.listTools().tools())
128131
.contains(Tool.builder().name("tool1").description("tool1 description").inputSchema("""
129132
{
@@ -166,6 +169,18 @@ void clientServerCapabilities() {
166169
.isEqualTo(json("""
167170
{"result":5.0,"operation":"2 + 3","timestamp":"2024-01-01T10:00:00Z"}"""));
168171

172+
// TOOL FROM MCP TOOL UTILS
173+
// Call the tool to ensure arguments are passed correctly
174+
CallToolResult toUpperCaseResponse = mcpClient
175+
.callTool(new McpSchema.CallToolRequest("toUpperCase", Map.of("input", "hello world")));
176+
assertThat(toUpperCaseResponse).isNotNull();
177+
assertThat(toUpperCaseResponse.isError()).isFalse();
178+
assertThat(toUpperCaseResponse.content()).hasSize(1)
179+
.first()
180+
.isInstanceOf(TextContent.class)
181+
.extracting("text")
182+
.isEqualTo("\"HELLO WORLD\"");
183+
169184
// PROMPT / COMPLETION
170185

171186
// list prompts
@@ -254,7 +269,20 @@ public List<McpStatelessServerFeatures.SyncToolSpecification> myTools() {
254269
})
255270
.build();
256271

257-
return List.of(tool1, tool2);
272+
// Tool 3
273+
274+
// Using a tool with McpToolUtils
275+
McpStatelessServerFeatures.SyncToolSpecification tool3 = McpToolUtils
276+
.toStatelessSyncToolSpecification(FunctionToolCallback
277+
.builder("toUpperCase", (ToUpperCaseRequest req, ToolContext context) -> req.input().toUpperCase())
278+
.description("Sets the input string to upper case")
279+
.inputType(ToUpperCaseRequest.class)
280+
.build(), null);
281+
282+
return List.of(tool1, tool2, tool3);
283+
}
284+
285+
record ToUpperCaseRequest(String input) {
258286
}
259287

260288
@Bean

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

Lines changed: 36 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import java.util.List;
2020
import java.util.Map;
2121
import java.util.Optional;
22+
import java.util.function.BiFunction;
2223

2324
import com.fasterxml.jackson.annotation.JsonAlias;
2425
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
@@ -29,7 +30,9 @@
2930
import io.modelcontextprotocol.server.McpServerFeatures.AsyncToolSpecification;
3031
import io.modelcontextprotocol.server.McpStatelessServerFeatures;
3132
import io.modelcontextprotocol.server.McpSyncServerExchange;
33+
import io.modelcontextprotocol.server.McpTransportContext;
3234
import io.modelcontextprotocol.spec.McpSchema;
35+
import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
3336
import io.modelcontextprotocol.spec.McpSchema.Role;
3437
import reactor.core.publisher.Mono;
3538
import reactor.core.scheduler.Schedulers;
@@ -152,16 +155,6 @@ public static McpServerFeatures.SyncToolSpecification toSyncToolSpecification(To
152155
* Converts a Spring AI ToolCallback to an MCP SyncToolSpecification. This enables
153156
* Spring AI functions to be exposed as MCP tools that can be discovered and invoked
154157
* by language models.
155-
*
156-
* <p>
157-
* The conversion process:
158-
* <ul>
159-
* <li>Creates an MCP Tool with the function's name and input schema</li>
160-
* <li>Wraps the function's execution in a SyncToolSpecification that handles the MCP
161-
* protocol</li>
162-
* <li>Provides error handling and result formatting according to MCP
163-
* specifications</li>
164-
* </ul>
165158
* @param toolCallback the Spring AI function callback to convert
166159
* @param mimeType the MIME type of the output content
167160
* @return an MCP SyncToolSpecification that wraps the function callback
@@ -170,39 +163,50 @@ public static McpServerFeatures.SyncToolSpecification toSyncToolSpecification(To
170163
public static McpServerFeatures.SyncToolSpecification toSyncToolSpecification(ToolCallback toolCallback,
171164
MimeType mimeType) {
172165

173-
var tool = new McpSchema.Tool(toolCallback.getToolDefinition().name(),
174-
toolCallback.getToolDefinition().description(), toolCallback.getToolDefinition().inputSchema());
166+
SharedSyncToolSpecification sharedSpec = toSharedSyncToolSpecification(toolCallback, mimeType);
175167

176-
return new McpServerFeatures.SyncToolSpecification(tool, (exchange, request) -> {
177-
try {
178-
String callResult = toolCallback.call(ModelOptionsUtils.toJsonString(request),
179-
new ToolContext(Map.of(TOOL_CONTEXT_MCP_EXCHANGE_KEY, exchange)));
180-
if (mimeType != null && mimeType.toString().startsWith("image")) {
181-
return new McpSchema.CallToolResult(List
182-
.of(new McpSchema.ImageContent(List.of(Role.ASSISTANT), null, callResult, mimeType.toString())),
183-
false);
184-
}
185-
return new McpSchema.CallToolResult(List.of(new McpSchema.TextContent(callResult)), false);
186-
}
187-
catch (Exception e) {
188-
return new McpSchema.CallToolResult(List.of(new McpSchema.TextContent(e.getMessage())), true);
189-
}
190-
});
168+
return new McpServerFeatures.SyncToolSpecification(sharedSpec.tool(), null,
169+
(exchange, request) -> sharedSpec.sharedHandler().apply(exchange, request));
191170
}
192171

172+
/**
173+
* Converts a Spring AI ToolCallback to an MCP StatelessSyncToolSpecification. This
174+
* enables Spring AI functions to be exposed as MCP tools that can be discovered and
175+
* invoked by language models.
176+
*
177+
* You can use the ToolCallback builder to create a new instance of ToolCallback using
178+
* either java.util.function.Function or Method reference.
179+
* @param toolCallback the Spring AI function callback to convert
180+
* @param mimeType the MIME type of the output content
181+
* @return an MCP StatelessSyncToolSpecification that wraps the function callback
182+
* @throws RuntimeException if there's an error during the function execution
183+
*/
193184
public static McpStatelessServerFeatures.SyncToolSpecification toStatelessSyncToolSpecification(
194185
ToolCallback toolCallback, MimeType mimeType) {
195186

187+
var sharedSpec = toSharedSyncToolSpecification(toolCallback, mimeType);
188+
189+
return new McpStatelessServerFeatures.SyncToolSpecification(sharedSpec.tool(),
190+
(exchange, request) -> sharedSpec.sharedHandler().apply(exchange, request));
191+
}
192+
193+
private record SharedSyncToolSpecification(McpSchema.Tool tool,
194+
BiFunction<Object, CallToolRequest, McpSchema.CallToolResult> sharedHandler) {
195+
}
196+
197+
private static SharedSyncToolSpecification toSharedSyncToolSpecification(ToolCallback toolCallback,
198+
MimeType mimeType) {
199+
196200
var tool = McpSchema.Tool.builder()
197201
.name(toolCallback.getToolDefinition().name())
198202
.description(toolCallback.getToolDefinition().description())
199203
.inputSchema(toolCallback.getToolDefinition().inputSchema())
200204
.build();
201205

202-
return new McpStatelessServerFeatures.SyncToolSpecification(tool, (mcpTransportContext, request) -> {
206+
return new SharedSyncToolSpecification(tool, (exchangeOrContext, request) -> {
203207
try {
204-
String callResult = toolCallback.call(ModelOptionsUtils.toJsonString(request),
205-
new ToolContext(Map.of(TOOL_CONTEXT_MCP_EXCHANGE_KEY, mcpTransportContext)));
208+
String callResult = toolCallback.call(ModelOptionsUtils.toJsonString(request.arguments()),
209+
new ToolContext(Map.of(TOOL_CONTEXT_MCP_EXCHANGE_KEY, exchangeOrContext)));
206210
if (mimeType != null && mimeType.toString().startsWith("image")) {
207211
return new McpSchema.CallToolResult(List
208212
.of(new McpSchema.ImageContent(List.of(Role.ASSISTANT), null, callResult, mimeType.toString())),
@@ -329,8 +333,8 @@ public static McpStatelessServerFeatures.AsyncToolSpecification toStatelessAsync
329333
toolCallback, mimeType);
330334

331335
return new McpStatelessServerFeatures.AsyncToolSpecification(statelessSyncToolSpecification.tool(),
332-
(context, map) -> Mono
333-
.fromCallable(() -> statelessSyncToolSpecification.callHandler().apply(context, map))
336+
(context, request) -> Mono
337+
.fromCallable(() -> statelessSyncToolSpecification.callHandler().apply(context.copy(), request))
334338
.subscribeOn(Schedulers.boundedElastic()));
335339
}
336340

0 commit comments

Comments
 (0)