Skip to content

Commit 5783104

Browse files
tzolovilayaperumalg
tzolov
authored and
ilayaperumalg
committed
feat(mcp): add tool callback caching with event-driven invalidation
- Add caching to SyncMcpToolCallbackProvider and AsyncMcpToolCallbackProvider to avoid re-discovering tools on every getToolCallbacks() call - Implement ApplicationListener<McpToolsChangedEvent> in both providers to invalidate cache when tools change - Add McpSyncToolsChangeEventEmmiter and McpAsyncToolsChangeEventEmmiter beans in McpClientAutoConfiguration for publishing tool change events - Use AtomicBoolean and double-checked locking pattern for thread-safe cache invalidation Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com> Use lock instead of synch blocks Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
1 parent a8d2208 commit 5783104

File tree

8 files changed

+467
-39
lines changed

8 files changed

+467
-39
lines changed

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
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.client.common.autoconfigure;
18+
19+
import io.modelcontextprotocol.client.McpClient.AsyncSpec;
20+
import io.modelcontextprotocol.util.Assert;
21+
import reactor.core.publisher.Mono;
22+
23+
import org.springframework.ai.mcp.McpToolsChangedEvent;
24+
import org.springframework.ai.mcp.customizer.McpAsyncClientCustomizer;
25+
import org.springframework.context.ApplicationEventPublisher;
26+
27+
/**
28+
* Emits {@link McpToolsChangedEvent} when the MCP Tools have changed for a given MCP
29+
* connection.
30+
*
31+
* @author Christian Tzolov
32+
*/
33+
public class McpAsyncToolsChangeEventEmmiter implements McpAsyncClientCustomizer {
34+
35+
private final ApplicationEventPublisher applicationEventPublisher;
36+
37+
public McpAsyncToolsChangeEventEmmiter(ApplicationEventPublisher applicationEventPublisher) {
38+
Assert.notNull(applicationEventPublisher, "applicationEventPublisher must not be null");
39+
this.applicationEventPublisher = applicationEventPublisher;
40+
}
41+
42+
@Override
43+
public void customize(String connectionName, AsyncSpec spec) {
44+
spec.toolsChangeConsumer(tools -> {
45+
this.applicationEventPublisher.publishEvent(new McpToolsChangedEvent(connectionName, tools));
46+
return Mono.empty();
47+
});
48+
}
49+
50+
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
3838
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3939
import org.springframework.boot.context.properties.EnableConfigurationProperties;
40+
import org.springframework.context.ApplicationEventPublisher;
4041
import org.springframework.context.annotation.Bean;
4142
import org.springframework.util.CollectionUtils;
4243

@@ -122,6 +123,14 @@ private String connectedClientName(String clientName, String serverConnectionNam
122123
return clientName + " - " + serverConnectionName;
123124
}
124125

126+
@Bean
127+
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "SYNC",
128+
matchIfMissing = true)
129+
public McpSyncToolsChangeEventEmmiter mcpSyncToolChangeEventEmmiter(
130+
ApplicationEventPublisher applicationEventPublisher) {
131+
return new McpSyncToolsChangeEventEmmiter(applicationEventPublisher);
132+
}
133+
125134
/**
126135
* Creates a list of {@link McpSyncClient} instances based on the available
127136
* transports.
@@ -226,6 +235,13 @@ McpSyncClientConfigurer mcpSyncClientConfigurer(ObjectProvider<McpSyncClientCust
226235

227236
// Async client configuration
228237

238+
@Bean
239+
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
240+
public McpAsyncToolsChangeEventEmmiter mcpAsyncToolChangeEventEmmiter(
241+
ApplicationEventPublisher applicationEventPublisher) {
242+
return new McpAsyncToolsChangeEventEmmiter(applicationEventPublisher);
243+
}
244+
229245
@Bean
230246
@ConditionalOnProperty(prefix = McpClientCommonProperties.CONFIG_PREFIX, name = "type", havingValue = "ASYNC")
231247
public List<McpAsyncClient> mcpAsyncClients(McpAsyncClientConfigurer mcpAsyncClientConfigurer,

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

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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.client.common.autoconfigure;
18+
19+
import io.modelcontextprotocol.client.McpClient.SyncSpec;
20+
import io.modelcontextprotocol.util.Assert;
21+
22+
import org.springframework.ai.mcp.McpToolsChangedEvent;
23+
import org.springframework.ai.mcp.customizer.McpSyncClientCustomizer;
24+
import org.springframework.context.ApplicationEventPublisher;
25+
26+
/**
27+
* Emits {@link McpToolsChangedEvent} when the MCP Tools have changed for a given MCP
28+
* connection.
29+
*
30+
* @author Christian Tzolov
31+
*/
32+
public class McpSyncToolsChangeEventEmmiter implements McpSyncClientCustomizer {
33+
34+
private final ApplicationEventPublisher applicationEventPublisher;
35+
36+
public McpSyncToolsChangeEventEmmiter(ApplicationEventPublisher applicationEventPublisher) {
37+
Assert.notNull(applicationEventPublisher, "applicationEventPublisher must not be null");
38+
this.applicationEventPublisher = applicationEventPublisher;
39+
}
40+
41+
@Override
42+
public void customize(String connectionName, SyncSpec spec) {
43+
spec.toolsChangeConsumer(
44+
tools -> this.applicationEventPublisher.publishEvent(new McpToolsChangedEvent(connectionName, tools)));
45+
46+
}
47+
48+
}

‎auto-configurations/mcp/spring-ai-autoconfigure-mcp-server-webflux/pom.xml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,12 @@
115115
<version>${project.parent.version}</version>
116116
<scope>test</scope>
117117
</dependency>
118+
119+
<dependency>
120+
<groupId>org.awaitility</groupId>
121+
<artifactId>awaitility</artifactId>
122+
<scope>test</scope>
123+
</dependency>
118124

119125
</dependencies>
120126

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

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
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.server.autoconfigure;
18+
19+
import java.time.Duration;
20+
import java.util.Arrays;
21+
import java.util.function.Function;
22+
import java.util.stream.Stream;
23+
24+
import io.modelcontextprotocol.server.McpSyncServer;
25+
import io.modelcontextprotocol.server.transport.WebFluxStreamableServerTransportProvider;
26+
import org.junit.jupiter.api.Test;
27+
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
28+
import org.springaicommunity.mcp.annotation.McpTool;
29+
import org.springaicommunity.mcp.annotation.McpToolParam;
30+
import org.springaicommunity.mcp.context.McpSyncRequestContext;
31+
import reactor.netty.DisposableServer;
32+
import reactor.netty.http.server.HttpServer;
33+
34+
import org.springframework.ai.mcp.McpToolUtils;
35+
import org.springframework.ai.mcp.client.common.autoconfigure.McpClientAutoConfiguration;
36+
import org.springframework.ai.mcp.client.common.autoconfigure.McpToolCallbackAutoConfiguration;
37+
import org.springframework.ai.mcp.client.common.autoconfigure.annotations.McpClientAnnotationScannerAutoConfiguration;
38+
import org.springframework.ai.mcp.client.webflux.autoconfigure.StreamableHttpWebFluxTransportAutoConfiguration;
39+
import org.springframework.ai.mcp.server.common.autoconfigure.McpServerAutoConfiguration;
40+
import org.springframework.ai.mcp.server.common.autoconfigure.McpServerObjectMapperAutoConfiguration;
41+
import org.springframework.ai.mcp.server.common.autoconfigure.ToolCallbackConverterAutoConfiguration;
42+
import org.springframework.ai.mcp.server.common.autoconfigure.annotations.McpServerAnnotationScannerAutoConfiguration;
43+
import org.springframework.ai.mcp.server.common.autoconfigure.annotations.McpServerSpecificationFactoryAutoConfiguration;
44+
import org.springframework.ai.model.anthropic.autoconfigure.AnthropicChatAutoConfiguration;
45+
import org.springframework.ai.model.chat.client.autoconfigure.ChatClientAutoConfiguration;
46+
import org.springframework.ai.model.tool.autoconfigure.ToolCallingAutoConfiguration;
47+
import org.springframework.ai.retry.autoconfigure.SpringAiRetryAutoConfiguration;
48+
import org.springframework.ai.tool.ToolCallbackProvider;
49+
import org.springframework.ai.tool.function.FunctionToolCallback;
50+
import org.springframework.boot.autoconfigure.AutoConfigurations;
51+
import org.springframework.boot.autoconfigure.web.client.RestClientAutoConfiguration;
52+
import org.springframework.boot.autoconfigure.web.reactive.function.client.WebClientAutoConfiguration;
53+
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
54+
import org.springframework.context.ApplicationContext;
55+
import org.springframework.context.annotation.Bean;
56+
import org.springframework.http.server.reactive.HttpHandler;
57+
import org.springframework.http.server.reactive.ReactorHttpHandlerAdapter;
58+
import org.springframework.test.util.TestSocketUtils;
59+
import org.springframework.web.reactive.function.server.RouterFunctions;
60+
61+
import static org.assertj.core.api.Assertions.assertThat;
62+
import static org.awaitility.Awaitility.await;
63+
64+
/**
65+
* @author Christian Tzolov
66+
*/
67+
@EnabledIfEnvironmentVariable(named = "ANTHROPIC_API_KEY", matches = ".*")
68+
public class McpToolCallProviderCachingIT {
69+
70+
private final ApplicationContextRunner serverContextRunner = new ApplicationContextRunner()
71+
.withPropertyValues("spring.ai.mcp.server.protocol=STREAMABLE")
72+
.withConfiguration(AutoConfigurations.of(McpServerAutoConfiguration.class,
73+
McpServerObjectMapperAutoConfiguration.class, ToolCallbackConverterAutoConfiguration.class,
74+
McpServerStreamableHttpWebFluxAutoConfiguration.class,
75+
McpServerAnnotationScannerAutoConfiguration.class,
76+
McpServerSpecificationFactoryAutoConfiguration.class));
77+
78+
private final ApplicationContextRunner clientApplicationContext = new ApplicationContextRunner()
79+
.withPropertyValues("spring.ai.anthropic.apiKey=" + System.getenv("ANTHROPIC_API_KEY"))
80+
.withConfiguration(anthropicAutoConfig(McpToolCallbackAutoConfiguration.class, McpClientAutoConfiguration.class,
81+
StreamableHttpWebFluxTransportAutoConfiguration.class,
82+
McpClientAnnotationScannerAutoConfiguration.class, AnthropicChatAutoConfiguration.class,
83+
ChatClientAutoConfiguration.class));
84+
85+
private static AutoConfigurations anthropicAutoConfig(Class<?>... additional) {
86+
Class<?>[] dependencies = { SpringAiRetryAutoConfiguration.class, ToolCallingAutoConfiguration.class,
87+
RestClientAutoConfiguration.class, WebClientAutoConfiguration.class };
88+
Class<?>[] all = Stream.concat(Arrays.stream(dependencies), Arrays.stream(additional)).toArray(Class<?>[]::new);
89+
return AutoConfigurations.of(all);
90+
}
91+
92+
@Test
93+
void clientToolCallbacksUpdateWhenServerToolsChangeAsync() {
94+
95+
int serverPort = TestSocketUtils.findAvailableTcpPort();
96+
97+
this.serverContextRunner.withUserConfiguration(TestMcpServerConfiguration.class)
98+
.withPropertyValues(// @formatter:off
99+
"spring.ai.mcp.server.name=test-mcp-server",
100+
"spring.ai.mcp.server.version=1.0.0",
101+
"spring.ai.mcp.server.streamable-http.keep-alive-interval=1s",
102+
"spring.ai.mcp.server.streamable-http.mcp-endpoint=/mcp") // @formatter:on
103+
.run(serverContext -> {
104+
105+
var httpServer = startHttpServer(serverContext, serverPort);
106+
107+
this.clientApplicationContext
108+
.withPropertyValues(// @formatter:off
109+
"spring.ai.mcp.client.streamable-http.connections.server1.url=http://localhost:" + serverPort,
110+
"spring.ai.mcp.client.initialized=false") // @formatter:on
111+
.run(clientContext -> {
112+
113+
ToolCallbackProvider tcp = clientContext.getBean(ToolCallbackProvider.class);
114+
115+
assertThat(tcp.getToolCallbacks()).hasSize(1);
116+
117+
McpSyncServer mcpSyncServer = serverContext.getBean(McpSyncServer.class);
118+
119+
var toolSpec = McpToolUtils
120+
.toSyncToolSpecification(FunctionToolCallback.builder("currentTime", new TimeService())
121+
.description("Get the current time by location")
122+
.inputType(TimeRequest.class)
123+
.build(), null);
124+
125+
mcpSyncServer.addTool(toolSpec);
126+
127+
// Wait for the tool to be added asynchronously
128+
await().atMost(Duration.ofSeconds(5))
129+
.pollInterval(Duration.ofMillis(100))
130+
.untilAsserted(() -> assertThat(tcp.getToolCallbacks()).hasSize(2));
131+
132+
mcpSyncServer.removeTool("weather");
133+
134+
// Wait for the tool to be removed asynchronously
135+
await().atMost(Duration.ofSeconds(5))
136+
.pollInterval(Duration.ofMillis(100))
137+
.untilAsserted(() -> assertThat(tcp.getToolCallbacks()).hasSize(1));
138+
});
139+
140+
stopHttpServer(httpServer);
141+
});
142+
}
143+
144+
// Helper methods to start and stop the HTTP server
145+
private static DisposableServer startHttpServer(ApplicationContext serverContext, int port) {
146+
WebFluxStreamableServerTransportProvider mcpStreamableServerTransport = serverContext
147+
.getBean(WebFluxStreamableServerTransportProvider.class);
148+
HttpHandler httpHandler = RouterFunctions.toHttpHandler(mcpStreamableServerTransport.getRouterFunction());
149+
ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler);
150+
return HttpServer.create().port(port).handle(adapter).bindNow();
151+
}
152+
153+
private static void stopHttpServer(DisposableServer server) {
154+
if (server != null) {
155+
server.disposeNow();
156+
}
157+
}
158+
159+
public static class TestMcpServerConfiguration {
160+
161+
@Bean
162+
public McpServerHandlers serverSideSpecProviders() {
163+
return new McpServerHandlers();
164+
}
165+
166+
public static class McpServerHandlers {
167+
168+
@McpTool(description = "Provides weather information by city name")
169+
public String weather(McpSyncRequestContext ctx, @McpToolParam String cityName) {
170+
return "Weather is 22C with rain ";
171+
}
172+
173+
}
174+
175+
}
176+
177+
public class TimeService implements Function<TimeRequest, String> {
178+
179+
public String apply(TimeRequest request) {
180+
return "The time in " + request.location() + " is 12:00 PM.";
181+
}
182+
183+
}
184+
185+
public record TimeRequest(String location) {
186+
}
187+
188+
}

0 commit comments

Comments
 (0)