Skip to content

Commit 7db214a

Browse files
committed
feat: MCP tools support ToolContext
- Added `excludedToolContextKeys` to `McpClientCommonProperties.ToolCallback`. This will let us control which `ToolContext` fields should not be sent to the server. By default, `excludedToolContextKeys` is set to `TOOL_CALL_HISTORY`. - Added `AbstractMcpToolCallback` to hold common code for `McpToolCallback` - Updated `AsyncMcpToolCallback` and `SyncMcpToolCallback` to pass `ToolContext` via `_meta` using key `ai.springframework.org/tool_context` - Updated `McpToolUtils` to help server tools extract `ToolContext` from MCP Client requests - Updated `AsyncMcpToolCallbackProvider` and `SyncMcpToolCallbackProvider` to use the Builder design pattern and add the `excludedToolContextKeys` field - Updated the corresponding test cases Signed-off-by: YunKui Lu <[email protected]>
1 parent 9d1e1b5 commit 7db214a

File tree

13 files changed

+658
-74
lines changed

13 files changed

+658
-74
lines changed

auto-configurations/mcp/spring-ai-autoconfigure-mcp-client-common/src/main/java/org/springframework/ai/mcp/client/common/autoconfigure/McpToolCallbackAutoConfiguration.java

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -51,16 +51,24 @@ public class McpToolCallbackAutoConfiguration {
5151
@Bean
5252
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "SYNC",
5353
matchIfMissing = true)
54-
public SyncMcpToolCallbackProvider mcpToolCallbacks(ObjectProvider<List<McpSyncClient>> syncMcpClients) {
54+
public SyncMcpToolCallbackProvider mcpToolCallbacks(ObjectProvider<List<McpSyncClient>> syncMcpClients,
55+
McpClientCommonProperties commonProperties) {
5556
List<McpSyncClient> mcpClients = syncMcpClients.stream().flatMap(List::stream).toList();
56-
return new SyncMcpToolCallbackProvider(mcpClients);
57+
return SyncMcpToolCallbackProvider.builder()
58+
.addMcpClients(mcpClients)
59+
.addExcludedToolContextKeys(commonProperties.getToolcallback().getExcludedToolContextKeys())
60+
.build();
5761
}
5862

5963
@Bean
6064
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
61-
public AsyncMcpToolCallbackProvider mcpAsyncToolCallbacks(ObjectProvider<List<McpAsyncClient>> mcpClientsProvider) {
65+
public AsyncMcpToolCallbackProvider mcpAsyncToolCallbacks(ObjectProvider<List<McpAsyncClient>> mcpClientsProvider,
66+
McpClientCommonProperties commonProperties) {
6267
List<McpAsyncClient> mcpClients = mcpClientsProvider.stream().flatMap(List::stream).toList();
63-
return new AsyncMcpToolCallbackProvider(mcpClients);
68+
return AsyncMcpToolCallbackProvider.builder()
69+
.addMcpClients(mcpClients)
70+
.addExcludedToolContextKeys(commonProperties.getToolcallback().getExcludedToolContextKeys())
71+
.build();
6472
}
6573

6674
public static class McpToolCallbackAutoConfigurationCondition extends AllNestedConditions {

auto-configurations/mcp/spring-ai-autoconfigure-mcp-client-common/src/main/java/org/springframework/ai/mcp/client/common/autoconfigure/properties/McpClientCommonProperties.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,12 @@
1717
package org.springframework.ai.mcp.client.common.autoconfigure.properties;
1818

1919
import java.time.Duration;
20+
import java.util.Set;
2021

2122
import org.springframework.boot.context.properties.ConfigurationProperties;
2223

24+
import static org.springframework.ai.chat.model.ToolContext.TOOL_CALL_HISTORY;
25+
2326
/**
2427
* Common Configuration properties for the Model Context Protocol (MCP) clients shared for
2528
* all transport types.
@@ -190,6 +193,14 @@ public static class Toolcallback {
190193
*/
191194
private boolean enabled = true;
192195

196+
/**
197+
* The keys that will not be sent to the MCP Server inside * the `_meta` field of
198+
* {@link io.modelcontextprotocol.spec.McpSchema.CallToolRequest}
199+
*/
200+
// Remember to update META-INF/additional-spring-configuration-metadata.json if
201+
// you change this default values.
202+
private Set<String> excludedToolContextKeys = Set.of(TOOL_CALL_HISTORY);
203+
193204
public void setEnabled(boolean enabled) {
194205
this.enabled = enabled;
195206
}
@@ -198,6 +209,14 @@ public boolean isEnabled() {
198209
return this.enabled;
199210
}
200211

212+
public Set<String> getExcludedToolContextKeys() {
213+
return excludedToolContextKeys;
214+
}
215+
216+
public void setExcludedToolContextKeys(Set<String> excludedToolContextKeys) {
217+
this.excludedToolContextKeys = excludedToolContextKeys;
218+
}
219+
201220
}
202221

203222
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{"properties": [
2+
{
3+
"name": "spring.ai.mcp.client.toolcallback.excluded-tool-context-keys",
4+
"defaultValue": ["TOOL_CALL_HISTORY"]
5+
}
6+
]}

auto-configurations/mcp/spring-ai-autoconfigure-mcp-client-common/src/test/java/org/springframework/ai/mcp/client/common/autoconfigure/properties/McpClientCommonPropertiesTests.java

Lines changed: 48 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package org.springframework.ai.mcp.client.common.autoconfigure.properties;
1818

1919
import java.time.Duration;
20+
import java.util.Set;
2021

2122
import org.junit.jupiter.api.Test;
2223

@@ -47,6 +48,9 @@ void defaultValues() {
4748
assertThat(properties.getRequestTimeout()).isEqualTo(Duration.ofSeconds(20));
4849
assertThat(properties.getType()).isEqualTo(McpClientCommonProperties.ClientType.SYNC);
4950
assertThat(properties.isRootChangeNotification()).isTrue();
51+
assertThat(properties.getToolcallback().isEnabled()).isTrue();
52+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
53+
.containsExactlyElementsOf(Set.of("TOOL_CALL_HISTORY"));
5054
});
5155
}
5256

@@ -56,7 +60,9 @@ void customValues() {
5660
.withPropertyValues("spring.ai.mcp.client.enabled=false", "spring.ai.mcp.client.name=custom-client",
5761
"spring.ai.mcp.client.version=2.0.0", "spring.ai.mcp.client.initialized=false",
5862
"spring.ai.mcp.client.request-timeout=30s", "spring.ai.mcp.client.type=ASYNC",
59-
"spring.ai.mcp.client.root-change-notification=false")
63+
"spring.ai.mcp.client.root-change-notification=false",
64+
"spring.ai.mcp.client.toolcallback.enabled=false",
65+
"spring.ai.mcp.client.toolcallback.excluded-tool-context-keys=foo,bar")
6066
.run(context -> {
6167
McpClientCommonProperties properties = context.getBean(McpClientCommonProperties.class);
6268
assertThat(properties.isEnabled()).isFalse();
@@ -66,6 +72,9 @@ void customValues() {
6672
assertThat(properties.getRequestTimeout()).isEqualTo(Duration.ofSeconds(30));
6773
assertThat(properties.getType()).isEqualTo(McpClientCommonProperties.ClientType.ASYNC);
6874
assertThat(properties.isRootChangeNotification()).isFalse();
75+
assertThat(properties.getToolcallback().isEnabled()).isFalse();
76+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
77+
.containsExactlyInAnyOrderElementsOf(Set.of("foo", "bar"));
6978
});
7079
}
7180

@@ -101,6 +110,14 @@ void setterGetterMethods() {
101110
// Test rootChangeNotification property
102111
properties.setRootChangeNotification(false);
103112
assertThat(properties.isRootChangeNotification()).isFalse();
113+
114+
// Test toolcallback property
115+
properties.getToolcallback().setEnabled(false);
116+
assertThat(properties.getToolcallback().isEnabled()).isFalse();
117+
118+
properties.getToolcallback().setExcludedToolContextKeys(Set.of("foo", "bar"));
119+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
120+
.containsExactlyInAnyOrderElementsOf(Set.of("foo", "bar"));
104121
}
105122

106123
@Test
@@ -125,7 +142,9 @@ void propertiesFileBinding() {
125142
.withPropertyValues("spring.ai.mcp.client.enabled=false", "spring.ai.mcp.client.name=test-mcp-client",
126143
"spring.ai.mcp.client.version=0.5.0", "spring.ai.mcp.client.initialized=false",
127144
"spring.ai.mcp.client.request-timeout=45s", "spring.ai.mcp.client.type=ASYNC",
128-
"spring.ai.mcp.client.root-change-notification=false")
145+
"spring.ai.mcp.client.root-change-notification=false",
146+
"spring.ai.mcp.client.toolcallback.enabled=false",
147+
"spring.ai.mcp.client.toolcallback.excluded-tool-context-keys=foo,bar")
129148
.run(context -> {
130149
McpClientCommonProperties properties = context.getBean(McpClientCommonProperties.class);
131150
assertThat(properties.isEnabled()).isFalse();
@@ -135,6 +154,9 @@ void propertiesFileBinding() {
135154
assertThat(properties.getRequestTimeout()).isEqualTo(Duration.ofSeconds(45));
136155
assertThat(properties.getType()).isEqualTo(McpClientCommonProperties.ClientType.ASYNC);
137156
assertThat(properties.isRootChangeNotification()).isFalse();
157+
assertThat(properties.getToolcallback().isEnabled()).isFalse();
158+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
159+
.containsExactlyInAnyOrderElementsOf(Set.of("foo", "bar"));
138160
});
139161
}
140162

@@ -165,7 +187,9 @@ void yamlConfigurationBinding() {
165187
.withPropertyValues("spring.ai.mcp.client.enabled=false", "spring.ai.mcp.client.name=test-mcp-client-yaml",
166188
"spring.ai.mcp.client.version=0.6.0", "spring.ai.mcp.client.initialized=false",
167189
"spring.ai.mcp.client.request-timeout=60s", "spring.ai.mcp.client.type=ASYNC",
168-
"spring.ai.mcp.client.root-change-notification=false")
190+
"spring.ai.mcp.client.root-change-notification=false",
191+
"spring.ai.mcp.client.toolcallback.enabled=false",
192+
"spring.ai.mcp.client.toolcallback.excluded-tool-context-keys=foo,bar")
169193
.run(context -> {
170194
McpClientCommonProperties properties = context.getBean(McpClientCommonProperties.class);
171195
assertThat(properties.isEnabled()).isFalse();
@@ -175,6 +199,9 @@ void yamlConfigurationBinding() {
175199
assertThat(properties.getRequestTimeout()).isEqualTo(Duration.ofSeconds(60));
176200
assertThat(properties.getType()).isEqualTo(McpClientCommonProperties.ClientType.ASYNC);
177201
assertThat(properties.isRootChangeNotification()).isFalse();
202+
assertThat(properties.getToolcallback().isEnabled()).isFalse();
203+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
204+
.containsExactlyInAnyOrderElementsOf(Set.of("foo", "bar"));
178205
});
179206
}
180207

@@ -201,6 +228,9 @@ void disabledProperties() {
201228
assertThat(properties.getRequestTimeout()).isEqualTo(Duration.ofSeconds(20));
202229
assertThat(properties.getType()).isEqualTo(McpClientCommonProperties.ClientType.SYNC);
203230
assertThat(properties.isRootChangeNotification()).isTrue();
231+
assertThat(properties.getToolcallback().isEnabled()).isTrue();
232+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
233+
.containsExactlyInAnyOrderElementsOf(Set.of("TOOL_CALL_HISTORY"));
204234
});
205235
}
206236

@@ -216,6 +246,9 @@ void notInitializedProperties() {
216246
assertThat(properties.getRequestTimeout()).isEqualTo(Duration.ofSeconds(20));
217247
assertThat(properties.getType()).isEqualTo(McpClientCommonProperties.ClientType.SYNC);
218248
assertThat(properties.isRootChangeNotification()).isTrue();
249+
assertThat(properties.getToolcallback().isEnabled()).isTrue();
250+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
251+
.containsExactlyInAnyOrderElementsOf(Set.of("TOOL_CALL_HISTORY"));
219252
});
220253
}
221254

@@ -231,6 +264,9 @@ void rootChangeNotificationDisabled() {
231264
assertThat(properties.isInitialized()).isTrue();
232265
assertThat(properties.getRequestTimeout()).isEqualTo(Duration.ofSeconds(20));
233266
assertThat(properties.getType()).isEqualTo(McpClientCommonProperties.ClientType.SYNC);
267+
assertThat(properties.getToolcallback().isEnabled()).isTrue();
268+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
269+
.containsExactlyInAnyOrderElementsOf(Set.of("TOOL_CALL_HISTORY"));
234270
});
235271
}
236272

@@ -246,6 +282,9 @@ void customRequestTimeout() {
246282
assertThat(properties.isInitialized()).isTrue();
247283
assertThat(properties.getType()).isEqualTo(McpClientCommonProperties.ClientType.SYNC);
248284
assertThat(properties.isRootChangeNotification()).isTrue();
285+
assertThat(properties.getToolcallback().isEnabled()).isTrue();
286+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
287+
.containsExactlyInAnyOrderElementsOf(Set.of("TOOL_CALL_HISTORY"));
249288
});
250289
}
251290

@@ -261,6 +300,9 @@ void asyncClientType() {
261300
assertThat(properties.isInitialized()).isTrue();
262301
assertThat(properties.getRequestTimeout()).isEqualTo(Duration.ofSeconds(20));
263302
assertThat(properties.isRootChangeNotification()).isTrue();
303+
assertThat(properties.getToolcallback().isEnabled()).isTrue();
304+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
305+
.containsExactlyInAnyOrderElementsOf(Set.of("TOOL_CALL_HISTORY"));
264306
});
265307
}
266308

@@ -278,6 +320,9 @@ void customNameAndVersion() {
278320
assertThat(properties.getRequestTimeout()).isEqualTo(Duration.ofSeconds(20));
279321
assertThat(properties.getType()).isEqualTo(McpClientCommonProperties.ClientType.SYNC);
280322
assertThat(properties.isRootChangeNotification()).isTrue();
323+
assertThat(properties.getToolcallback().isEnabled()).isTrue();
324+
assertThat(properties.getToolcallback().getExcludedToolContextKeys())
325+
.containsExactlyInAnyOrderElementsOf(Set.of("TOOL_CALL_HISTORY"));
281326
});
282327
}
283328

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
/*
2+
* Copyright 2025-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.mcp;
18+
19+
import java.util.HashMap;
20+
import java.util.Map;
21+
import java.util.Set;
22+
23+
import org.springframework.ai.chat.model.ToolContext;
24+
import org.springframework.ai.tool.ToolCallback;
25+
import org.springframework.lang.Nullable;
26+
import org.springframework.util.Assert;
27+
28+
import static io.modelcontextprotocol.spec.McpSchema.Tool;
29+
30+
/**
31+
* Implementation of {@link ToolCallback} that adapts MCP tools to Spring AI's tool
32+
* interface with asynchronous execution support.
33+
* <p>
34+
* This class acts as a bridge between the Model Context Protocol (MCP) and Spring AI's
35+
* tool system, allowing MCP tools to be used seamlessly within Spring AI applications.
36+
* It:
37+
* <ul>
38+
* <li>Converts MCP tool definitions to Spring AI tool definitions</li>
39+
* <li>Handles the asynchronous execution of tool calls through the MCP client</li>
40+
* <li>Manages JSON serialization/deserialization of tool inputs and outputs</li>
41+
* </ul>
42+
* <p>
43+
*
44+
* @author YunKui Lu
45+
* @see ToolCallback
46+
* @see AsyncMcpToolCallback
47+
* @see SyncMcpToolCallback
48+
* @see Tool
49+
*/
50+
public abstract class AbstractMcpToolCallback implements ToolCallback {
51+
52+
public static final String DEFAULT_MCP_META_TOOL_CONTEXT_KEY = McpToolUtils.DEFAULT_MCP_META_TOOL_CONTEXT_KEY;
53+
54+
protected final Tool tool;
55+
56+
/**
57+
* the keys that will not be sent to the MCP Server inside the `_meta` field of
58+
* {@link io.modelcontextprotocol.spec.McpSchema.CallToolRequest}
59+
*/
60+
protected final Set<String> excludedToolContextKeys;
61+
62+
/**
63+
* Creates a new {@code AbstractMcpToolCallback} instance.
64+
* @param tool the MCP tool definition to adapt
65+
* @param excludedToolContextKeys the keys that will not be sent to the MCP Server
66+
* inside the `_meta` field of
67+
* {@link io.modelcontextprotocol.spec.McpSchema.CallToolRequest}
68+
*/
69+
protected AbstractMcpToolCallback(Tool tool, Set<String> excludedToolContextKeys) {
70+
Assert.notNull(tool, "tool cannot be null");
71+
Assert.notNull(excludedToolContextKeys, "excludedToolContextKeys cannot be null");
72+
73+
this.tool = tool;
74+
this.excludedToolContextKeys = excludedToolContextKeys;
75+
}
76+
77+
/**
78+
* Executes the tool with the provided input.
79+
* <p>
80+
* This method:
81+
* <ol>
82+
* <li>Converts the JSON input string to a map of arguments</li>
83+
* <li>Calls the tool through the MCP client</li>
84+
* <li>Converts the tool's response content to a JSON string</li>
85+
* </ol>
86+
* @param functionInput the tool input as a JSON string
87+
* @return the tool's response as a JSON string
88+
*/
89+
@Override
90+
public String call(String functionInput) {
91+
return this.call(functionInput, null);
92+
}
93+
94+
/**
95+
* Converts the tool context to a mcp meta map
96+
* @param toolContext the context for tool execution in a function calling scenario
97+
* @return the mcp meta map
98+
*/
99+
protected Map<String, Object> getAdditionalToolContextToMeta(@Nullable ToolContext toolContext) {
100+
if (toolContext == null || toolContext.getContext().isEmpty()) {
101+
return Map.of();
102+
}
103+
104+
Map<String, Object> meta = new HashMap<>(toolContext.getContext().size() - excludedToolContextKeys.size());
105+
for (var toolContextEntry : toolContext.getContext().entrySet()) {
106+
if (excludedToolContextKeys.contains(toolContextEntry.getKey())) {
107+
continue;
108+
}
109+
meta.put(toolContextEntry.getKey(), toolContextEntry.getValue());
110+
}
111+
return meta;
112+
}
113+
114+
}

0 commit comments

Comments
 (0)