Skip to content

Commit 1fdda61

Browse files
tzolovtzolov
committed
feat(autoconfigure): Support both FunctionCallback and ToolCallback in ToolCallingAutoConfiguration
- Extends the ToolCallingAutoConfiguration to support both FunctionCallback and ToolCallback types. - The toolCallbackResolver bean now handles both callback types through ObjectProvider injection. - Added comprehensive tests to verify the resolution of multiple function and tool callbacks. - Introduce new StaticToolCallbackProvider implementation - Update ToolCallbackProvider to return FunctionCallback[] - Migrate from List to ToolCallbackProvider in configurations - Update tests to use new provider pattern - Enhance tool callback providers to support multiple clients - Refactor AsyncMcpToolCallbackProvider and SyncMcpToolCallbackProvider to handle multiple MCP clients - Add ToolCallbackProvider support to ChatClient API - Deprecate direct tool callback list methods in favor of providers - Fix typos in Closeable class names - Update MCP documentation with new examples and usage patterns Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
1 parent 68ad742 commit 1fdda61

File tree

16 files changed

+487
-109
lines changed

16 files changed

+487
-109
lines changed

‎auto-configurations/spring-ai-mcp-client/src/main/java/org/springframework/ai/autoconfigure/mcp/client/McpClientAutoConfiguration.java

Lines changed: 38 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,12 @@
2727
import org.springframework.ai.autoconfigure.mcp.client.configurer.McpAsyncClientConfigurer;
2828
import org.springframework.ai.autoconfigure.mcp.client.configurer.McpSyncClientConfigurer;
2929
import org.springframework.ai.autoconfigure.mcp.client.properties.McpClientCommonProperties;
30-
import org.springframework.ai.mcp.McpToolUtils;
30+
import org.springframework.ai.mcp.AsyncMcpToolCallbackProvider;
31+
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
3132
import org.springframework.ai.mcp.customizer.McpAsyncClientCustomizer;
3233
import org.springframework.ai.mcp.customizer.McpSyncClientCustomizer;
3334
import org.springframework.ai.tool.ToolCallback;
35+
import org.springframework.ai.tool.ToolCallbackProvider;
3436
import org.springframework.beans.factory.ObjectProvider;
3537
import org.springframework.boot.autoconfigure.AutoConfiguration;
3638
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -176,9 +178,22 @@ public List<McpSyncClient> mcpSyncClients(McpSyncClientConfigurer mcpSyncClientC
176178
@Bean
177179
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "SYNC",
178180
matchIfMissing = true)
179-
public List<ToolCallback> toolCallbacks(ObjectProvider<List<McpSyncClient>> mcpClientsProvider) {
181+
public ToolCallbackProvider toolCallbacks(ObjectProvider<List<McpSyncClient>> mcpClientsProvider) {
180182
List<McpSyncClient> mcpClients = mcpClientsProvider.stream().flatMap(List::stream).toList();
181-
return McpToolUtils.getToolCallbacksFromSyncClients(mcpClients);
183+
return new SyncMcpToolCallbackProvider(mcpClients);
184+
}
185+
186+
/**
187+
* @deprecated replaced by {@link #toolCallbacks(ObjectProvider)} that returns a
188+
* {@link ToolCallbackProvider} instead of a list of {@link ToolCallback}
189+
*/
190+
@Deprecated
191+
@Bean
192+
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "SYNC",
193+
matchIfMissing = true)
194+
public List<ToolCallback> toolCallbacksDeprecated(ObjectProvider<List<McpSyncClient>> mcpClientsProvider) {
195+
List<McpSyncClient> mcpClients = mcpClientsProvider.stream().flatMap(List::stream).toList();
196+
return List.of(new SyncMcpToolCallbackProvider(mcpClients).getToolCallbacks());
182197
}
183198

184199
/**
@@ -189,7 +204,7 @@ public List<ToolCallback> toolCallbacks(ObjectProvider<List<McpSyncClient>> mcpC
189204
* This class is responsible for closing all MCP sync clients when the application
190205
* context is closed, preventing resource leaks.
191206
*/
192-
public record ClosebleMcpSyncClients(List<McpSyncClient> clients) implements AutoCloseable {
207+
public record CloseableMcpSyncClients(List<McpSyncClient> clients) implements AutoCloseable {
193208

194209
@Override
195210
public void close() {
@@ -205,8 +220,8 @@ public void close() {
205220
@Bean
206221
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "SYNC",
207222
matchIfMissing = true)
208-
public ClosebleMcpSyncClients makeSyncClientsClosable(List<McpSyncClient> clients) {
209-
return new ClosebleMcpSyncClients(clients);
223+
public CloseableMcpSyncClients makeSyncClientsClosable(List<McpSyncClient> clients) {
224+
return new CloseableMcpSyncClients(clients);
210225
}
211226

212227
/**
@@ -263,14 +278,26 @@ public List<McpAsyncClient> mcpAsyncClients(McpAsyncClientConfigurer mcpSyncClie
263278
return mcpSyncClients;
264279
}
265280

281+
/**
282+
* @deprecated replaced by {@link #asyncToolCallbacks(ObjectProvider)} that returns a
283+
* {@link ToolCallbackProvider} instead of a list of {@link ToolCallback}
284+
*/
285+
@Deprecated
286+
@Bean
287+
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
288+
public List<ToolCallback> asyncToolCallbacksDeprecated(ObjectProvider<List<McpAsyncClient>> mcpClientsProvider) {
289+
List<McpAsyncClient> mcpClients = mcpClientsProvider.stream().flatMap(List::stream).toList();
290+
return List.of(new AsyncMcpToolCallbackProvider(mcpClients).getToolCallbacks());
291+
}
292+
266293
@Bean
267294
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
268-
public List<ToolCallback> asyncToolCallbacks(ObjectProvider<List<McpAsyncClient>> mcpClientsProvider) {
295+
public ToolCallbackProvider asyncToolCallbacks(ObjectProvider<List<McpAsyncClient>> mcpClientsProvider) {
269296
List<McpAsyncClient> mcpClients = mcpClientsProvider.stream().flatMap(List::stream).toList();
270-
return McpToolUtils.getToolCallbacksFromAsyncClinents(mcpClients);
297+
return new AsyncMcpToolCallbackProvider(mcpClients);
271298
}
272299

273-
public record ClosebleMcpAsyncClients(List<McpAsyncClient> clients) implements AutoCloseable {
300+
public record CloseableMcpAsyncClients(List<McpAsyncClient> clients) implements AutoCloseable {
274301
@Override
275302
public void close() {
276303
this.clients.forEach(McpAsyncClient::close);
@@ -279,8 +306,8 @@ public void close() {
279306

280307
@Bean
281308
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
282-
public ClosebleMcpAsyncClients makeAsynClientsClosable(List<McpAsyncClient> clients) {
283-
return new ClosebleMcpAsyncClients(clients);
309+
public CloseableMcpAsyncClients makeAsynClientsClosable(List<McpAsyncClient> clients) {
310+
return new CloseableMcpAsyncClients(clients);
284311
}
285312

286313
@Bean

‎auto-configurations/spring-ai-mcp-client/src/test/java/org/springframework/ai/autoconfigure/mcp/client/McpClientAutoConfigurationIT.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ void toolCallbacksCreation() {
122122
@Test
123123
void closeableWrappersCreation() {
124124
this.contextRunner.withUserConfiguration(TestTransportConfiguration.class).run(context -> {
125-
assertThat(context).hasSingleBean(McpClientAutoConfiguration.ClosebleMcpSyncClients.class);
125+
assertThat(context).hasSingleBean(McpClientAutoConfiguration.CloseableMcpSyncClients.class);
126126
});
127127
}
128128

‎auto-configurations/spring-ai-mcp-server/src/main/java/org/springframework/ai/autoconfigure/mcp/server/MpcServerAutoConfiguration.java

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616

1717
package org.springframework.ai.autoconfigure.mcp.server;
1818

19+
import java.util.ArrayList;
1920
import java.util.List;
2021
import java.util.function.Consumer;
2122
import java.util.function.Function;
23+
import java.util.stream.Stream;
2224

2325
import io.modelcontextprotocol.server.McpAsyncServer;
2426
import io.modelcontextprotocol.server.McpServer;
@@ -39,7 +41,9 @@
3941
import reactor.core.publisher.Mono;
4042

4143
import org.springframework.ai.mcp.McpToolUtils;
44+
import org.springframework.ai.model.function.FunctionCallback;
4245
import org.springframework.ai.tool.ToolCallback;
46+
import org.springframework.ai.tool.ToolCallbackProvider;
4347
import org.springframework.beans.factory.ObjectProvider;
4448
import org.springframework.boot.autoconfigure.AutoConfiguration;
4549
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -135,15 +139,23 @@ public McpSyncServer mcpSyncServer(ServerMcpTransport transport,
135139
McpSchema.ServerCapabilities.Builder capabilitiesBuilder, McpServerProperties serverProperties,
136140
ObjectProvider<List<SyncToolRegistration>> tools, ObjectProvider<List<SyncResourceRegistration>> resources,
137141
ObjectProvider<List<SyncPromptRegistration>> prompts,
138-
ObjectProvider<Consumer<List<McpSchema.Root>>> rootsChangeConsumers) {
142+
ObjectProvider<Consumer<List<McpSchema.Root>>> rootsChangeConsumers,
143+
List<ToolCallbackProvider> toolCallbackProvider) {
139144

140145
McpSchema.Implementation serverInfo = new Implementation(serverProperties.getName(),
141146
serverProperties.getVersion());
142147

143148
// Create the server with both tool and resource capabilities
144149
SyncSpec serverBuilder = McpServer.sync(transport).serverInfo(serverInfo);
145150

146-
List<SyncToolRegistration> toolResgistrations = tools.stream().flatMap(List::stream).toList();
151+
List<SyncToolRegistration> toolResgistrations = new ArrayList<>(tools.stream().flatMap(List::stream).toList());
152+
List<ToolCallback> providerToolCallbacks = toolCallbackProvider.stream()
153+
.map(pr -> List.of(pr.getToolCallbacks()))
154+
.flatMap(List::stream)
155+
.filter(fc -> fc instanceof ToolCallback)
156+
.map(fc -> (ToolCallback) fc)
157+
.toList();
158+
toolResgistrations.addAll(McpToolUtils.toSyncToolRegistration(providerToolCallbacks));
147159
if (!CollectionUtils.isEmpty(toolResgistrations)) {
148160
serverBuilder.tools(toolResgistrations);
149161
capabilitiesBuilder.tools(serverProperties.isToolChangeNotification());
@@ -191,15 +203,23 @@ public McpAsyncServer mcpAsyncServer(ServerMcpTransport transport,
191203
ObjectProvider<List<AsyncToolRegistration>> tools,
192204
ObjectProvider<List<AsyncResourceRegistration>> resources,
193205
ObjectProvider<List<AsyncPromptRegistration>> prompts,
194-
ObjectProvider<Consumer<List<McpSchema.Root>>> rootsChangeConsumer) {
206+
ObjectProvider<Consumer<List<McpSchema.Root>>> rootsChangeConsumer,
207+
List<ToolCallbackProvider> toolCallbackProvider) {
195208

196209
McpSchema.Implementation serverInfo = new Implementation(serverProperties.getName(),
197210
serverProperties.getVersion());
198211

199212
// Create the server with both tool and resource capabilities
200213
AsyncSpec serverBilder = McpServer.async(transport).serverInfo(serverInfo);
201214

202-
List<AsyncToolRegistration> toolResgistrations = tools.stream().flatMap(List::stream).toList();
215+
List<AsyncToolRegistration> toolResgistrations = new ArrayList<>(tools.stream().flatMap(List::stream).toList());
216+
List<ToolCallback> providerToolCallbacks = toolCallbackProvider.stream()
217+
.map(pr -> List.of(pr.getToolCallbacks()))
218+
.flatMap(List::stream)
219+
.filter(fc -> fc instanceof ToolCallback)
220+
.map(fc -> (ToolCallback) fc)
221+
.toList();
222+
toolResgistrations.addAll(McpToolUtils.toAsyncToolRegistration(providerToolCallbacks));
203223
if (!CollectionUtils.isEmpty(toolResgistrations)) {
204224
serverBilder.tools(toolResgistrations);
205225
capabilitiesBuilder.tools(serverProperties.isToolChangeNotification());

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

Lines changed: 73 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,12 @@
1515
*/
1616
package org.springframework.ai.mcp;
1717

18+
import java.util.ArrayList;
1819
import java.util.List;
1920

2021
import io.modelcontextprotocol.client.McpAsyncClient;
22+
import io.modelcontextprotocol.util.Assert;
2123
import reactor.core.publisher.Flux;
22-
import reactor.core.publisher.Mono;
2324

2425
import org.springframework.ai.tool.ToolCallback;
2526
import org.springframework.ai.tool.ToolCallbackProvider;
@@ -28,25 +29,40 @@
2829

2930
/**
3031
* Implementation of {@link ToolCallbackProvider} that discovers and provides MCP tools
31-
* asynchronously.
32+
* asynchronously from one or more MCP servers.
3233
* <p>
3334
* This class acts as a tool provider for Spring AI, automatically discovering tools from
34-
* an MCP server and making them available as Spring AI tools. It:
35+
* multiple MCP servers and making them available as Spring AI tools. It:
3536
* <ul>
36-
* <li>Connects to an MCP server through an async client</li>
37-
* <li>Lists and retrieves available tools from the server</li>
37+
* <li>Connects to MCP servers through async clients</li>
38+
* <li>Lists and retrieves available tools from each server asynchronously</li>
3839
* <li>Creates {@link AsyncMcpToolCallback} instances for each discovered tool</li>
39-
* <li>Validates tool names to prevent duplicates</li>
40+
* <li>Validates tool names to prevent duplicates across all servers</li>
4041
* </ul>
4142
* <p>
42-
* Example usage: <pre>{@code
43+
* Example usage with a single client:
44+
*
45+
* <pre>{@code
4346
* McpAsyncClient mcpClient = // obtain MCP client
4447
* ToolCallbackProvider provider = new AsyncMcpToolCallbackProvider(mcpClient);
4548
*
4649
* // Get all available tools
4750
* ToolCallback[] tools = provider.getToolCallbacks();
4851
* }</pre>
4952
*
53+
* Example usage with multiple clients:
54+
*
55+
* <pre>{@code
56+
* List<McpAsyncClient> mcpClients = // obtain multiple MCP clients
57+
* ToolCallbackProvider provider = new AsyncMcpToolCallbackProvider(mcpClients);
58+
*
59+
* // Get tools from all clients
60+
* ToolCallback[] tools = provider.getToolCallbacks();
61+
*
62+
* // Or use the reactive API
63+
* Flux<ToolCallback> toolsFlux = AsyncMcpToolCallbackProvider.asyncToolCallbacks(mcpClients);
64+
* }</pre>
65+
*
5066
* @author Christian Tzolov
5167
* @since 1.0.0
5268
* @see ToolCallbackProvider
@@ -55,40 +71,61 @@
5571
*/
5672
public class AsyncMcpToolCallbackProvider implements ToolCallbackProvider {
5773

58-
private final McpAsyncClient mcpClient;
74+
private final List<McpAsyncClient> mcpClients;
5975

6076
/**
61-
* Creates a new {@code AsyncMcpToolCallbackProvider} instance.
62-
* @param mcpClient the MCP client to use for discovering tools
77+
* Creates a new {@code AsyncMcpToolCallbackProvider} instance with a list of MCP
78+
* clients.
79+
* @param mcpClients the list of MCP clients to use for discovering tools. Each client
80+
* typically connects to a different MCP server, allowing tool discovery from multiple
81+
* sources.
82+
* @throws IllegalArgumentException if mcpClients is null
6383
*/
64-
public AsyncMcpToolCallbackProvider(McpAsyncClient mcpClient) {
65-
this.mcpClient = mcpClient;
84+
public AsyncMcpToolCallbackProvider(List<McpAsyncClient> mcpClients) {
85+
Assert.notNull(mcpClients, "McpClients must not be null");
86+
this.mcpClients = mcpClients;
87+
}
88+
89+
public AsyncMcpToolCallbackProvider(McpAsyncClient... mcpClients) {
90+
Assert.notNull(mcpClients, "McpClients must not be null");
91+
this.mcpClients = List.of(mcpClients);
6692
}
6793

6894
/**
69-
* Discovers and returns all available tools from the MCP server asynchronously.
95+
* Discovers and returns all available tools from the configured MCP servers.
7096
* <p>
7197
* This method:
7298
* <ol>
73-
* <li>Retrieves the list of tools from the MCP server</li>
74-
* <li>Creates a {@link AsyncMcpToolCallback} for each tool</li>
75-
* <li>Validates that there are no duplicate tool names</li>
99+
* <li>Retrieves the list of tools from each MCP server asynchronously</li>
100+
* <li>Creates a {@link AsyncMcpToolCallback} for each discovered tool</li>
101+
* <li>Validates that there are no duplicate tool names across all servers</li>
76102
* </ol>
103+
* <p>
104+
* Note: While the underlying tool discovery is asynchronous, this method blocks until
105+
* all tools are discovered from all servers.
77106
* @return an array of tool callbacks, one for each discovered tool
78107
* @throws IllegalStateException if duplicate tool names are found
79108
*/
80109
@Override
81110
public ToolCallback[] getToolCallbacks() {
82-
var toolCallbacks = this.mcpClient.listTools()
83-
.map(response -> response.tools()
84-
.stream()
85-
.map(tool -> new AsyncMcpToolCallback(this.mcpClient, tool))
86-
.toArray(ToolCallback[]::new))
87-
.block();
88111

89-
validateToolCallbacks(toolCallbacks);
112+
List<ToolCallback> toolCallbackList = new ArrayList<>();
113+
114+
for (McpAsyncClient mcpClient : this.mcpClients) {
115+
116+
ToolCallback[] toolCallbacks = mcpClient.listTools()
117+
.map(response -> response.tools()
118+
.stream()
119+
.map(tool -> new AsyncMcpToolCallback(mcpClient, tool))
120+
.toArray(ToolCallback[]::new))
121+
.block();
90122

91-
return toolCallbacks;
123+
validateToolCallbacks(toolCallbacks);
124+
125+
toolCallbackList.addAll(List.of(toolCallbacks));
126+
}
127+
128+
return toolCallbackList.toArray(new ToolCallback[0]);
92129
}
93130

94131
/**
@@ -110,12 +147,19 @@ private void validateToolCallbacks(ToolCallback[] toolCallbacks) {
110147
/**
111148
* Creates a reactive stream of tool callbacks from multiple MCP clients.
112149
* <p>
113-
* This utility method:
150+
* This utility method provides a reactive way to work with tool callbacks from
151+
* multiple MCP clients in a single operation. It:
114152
* <ol>
115-
* <li>Takes a list of MCP clients</li>
116-
* <li>Creates a provider for each client</li>
117-
* <li>Retrieves and flattens all tool callbacks into a single stream</li>
153+
* <li>Takes a list of MCP clients as input</li>
154+
* <li>Creates a provider instance to manage all clients</li>
155+
* <li>Retrieves tools from all clients asynchronously</li>
156+
* <li>Combines them into a single reactive stream</li>
157+
* <li>Ensures there are no naming conflicts between tools from different clients</li>
118158
* </ol>
159+
* <p>
160+
* Unlike {@link #getToolCallbacks()}, this method provides a fully reactive way to
161+
* work with tool callbacks, making it suitable for non-blocking applications. Any
162+
* errors during tool discovery will be propagated through the returned Flux.
119163
* @param mcpClients the list of MCP clients to create callbacks from
120164
* @return a Flux of tool callbacks from all provided clients
121165
*/
@@ -124,9 +168,7 @@ public static Flux<ToolCallback> asyncToolCallbacks(List<McpAsyncClient> mcpClie
124168
return Flux.empty();
125169
}
126170

127-
return Flux.fromIterable(mcpClients)
128-
.flatMap(mcpClient -> Mono.just(new AsyncMcpToolCallbackProvider(mcpClient).getToolCallbacks()))
129-
.flatMap(callbacks -> Flux.fromArray(callbacks));
171+
return Flux.fromArray(new AsyncMcpToolCallbackProvider(mcpClients).getToolCallbacks());
130172
}
131173

132174
}

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -211,10 +211,7 @@ public static List<ToolCallback> getToolCallbacksFromSyncClients(List<McpSyncCli
211211
if (CollectionUtils.isEmpty(mcpClients)) {
212212
return List.of();
213213
}
214-
return mcpClients.stream()
215-
.map(mcpClient -> List.of((new SyncMcpToolCallbackProvider(mcpClient).getToolCallbacks())))
216-
.flatMap(List::stream)
217-
.toList();
214+
return List.of((new SyncMcpToolCallbackProvider(mcpClients).getToolCallbacks()));
218215
}
219216

220217
/**
@@ -247,10 +244,7 @@ public static List<ToolCallback> getToolCallbacksFromAsyncClinents(List<McpAsync
247244
if (CollectionUtils.isEmpty(asynMcpClients)) {
248245
return List.of();
249246
}
250-
return asynMcpClients.stream()
251-
.map(mcpClient -> List.of((new AsyncMcpToolCallbackProvider(mcpClient).getToolCallbacks())))
252-
.flatMap(List::stream)
253-
.toList();
247+
return List.of((new AsyncMcpToolCallbackProvider(asynMcpClients).getToolCallbacks()));
254248
}
255249

256250
}

0 commit comments

Comments
 (0)