Skip to content

Commit 662d2a3

Browse files
tzolovjoshlong
authored and
Willam2004
committed
adding AOT support for MCP annotations (spring-projects#4427)
* wip * works * up * Fix checkstyle issues Signed-off-by: Christian Tzolov <[email protected]> * checkstyle fixes Signed-off-by: Christian Tzolov <[email protected]> * remove sysout Signed-off-by: Christian Tzolov <[email protected]> --------- Signed-off-by: Christian Tzolov <[email protected]> Co-authored-by: Josh Long <[email protected]> Signed-off-by: 家娃 <[email protected]>
1 parent f3ea9b8 commit 662d2a3

File tree

15 files changed

+244
-86
lines changed

15 files changed

+244
-86
lines changed

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,23 +24,30 @@
2424
import org.springaicommunity.mcp.annotation.McpProgress;
2525
import org.springaicommunity.mcp.annotation.McpSampling;
2626

27+
import org.springframework.ai.mcp.annotation.spring.scan.AbstractAnnotatedMethodBeanFactoryInitializationAotProcessor;
2728
import org.springframework.ai.mcp.annotation.spring.scan.AbstractAnnotatedMethodBeanPostProcessor;
2829
import org.springframework.ai.mcp.annotation.spring.scan.AbstractMcpAnnotatedBeans;
30+
import org.springframework.aot.hint.MemberCategory;
31+
import org.springframework.aot.hint.RuntimeHints;
32+
import org.springframework.aot.hint.RuntimeHintsRegistrar;
2933
import org.springframework.boot.autoconfigure.AutoConfiguration;
3034
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
3135
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
3236
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3337
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3438
import org.springframework.context.annotation.Bean;
39+
import org.springframework.context.annotation.ImportRuntimeHints;
3540

3641
/**
3742
* @author Christian Tzolov
43+
* @author Josh Long
3844
*/
3945
@AutoConfiguration
4046
@ConditionalOnClass(McpLogging.class)
4147
@ConditionalOnProperty(prefix = McpClientAnnotationScannerProperties.CONFIG_PREFIX, name = "enabled",
4248
havingValue = "true", matchIfMissing = true)
4349
@EnableConfigurationProperties(McpClientAnnotationScannerProperties.class)
50+
@ImportRuntimeHints(McpClientAnnotationScannerAutoConfiguration.AnnotationHints.class)
4451
public class McpClientAnnotationScannerAutoConfiguration {
4552

4653
private static final Set<Class<? extends Annotation>> CLIENT_MCP_ANNOTATIONS = Set.of(McpLogging.class,
@@ -54,15 +61,30 @@ public ClientMcpAnnotatedBeans clientAnnotatedBeans() {
5461

5562
@Bean
5663
@ConditionalOnMissingBean
57-
public ClientAnnotatedMethodBeanPostProcessor clientAnnotatedMethodBeanPostProcessor(
64+
public static ClientAnnotatedMethodBeanPostProcessor clientAnnotatedMethodBeanPostProcessor(
5865
ClientMcpAnnotatedBeans clientMcpAnnotatedBeans, McpClientAnnotationScannerProperties properties) {
5966
return new ClientAnnotatedMethodBeanPostProcessor(clientMcpAnnotatedBeans, CLIENT_MCP_ANNOTATIONS);
6067
}
6168

69+
@Bean
70+
static ClientAnnotatedBeanFactoryInitializationAotProcessor clientAnnotatedBeanFactoryInitializationAotProcessor() {
71+
return new ClientAnnotatedBeanFactoryInitializationAotProcessor(CLIENT_MCP_ANNOTATIONS);
72+
}
73+
6274
public static class ClientMcpAnnotatedBeans extends AbstractMcpAnnotatedBeans {
6375

6476
}
6577

78+
public static class ClientAnnotatedBeanFactoryInitializationAotProcessor
79+
extends AbstractAnnotatedMethodBeanFactoryInitializationAotProcessor {
80+
81+
public ClientAnnotatedBeanFactoryInitializationAotProcessor(
82+
Set<Class<? extends Annotation>> targetAnnotations) {
83+
super(targetAnnotations);
84+
}
85+
86+
}
87+
6688
public static class ClientAnnotatedMethodBeanPostProcessor extends AbstractAnnotatedMethodBeanPostProcessor {
6789

6890
public ClientAnnotatedMethodBeanPostProcessor(ClientMcpAnnotatedBeans clientMcpAnnotatedBeans,
@@ -72,4 +94,13 @@ public ClientAnnotatedMethodBeanPostProcessor(ClientMcpAnnotatedBeans clientMcpA
7294

7395
}
7496

97+
static class AnnotationHints implements RuntimeHintsRegistrar {
98+
99+
@Override
100+
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
101+
CLIENT_MCP_ANNOTATIONS.forEach(an -> hints.reflection().registerType(an, MemberCategory.values()));
102+
}
103+
104+
}
105+
75106
}

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,8 +94,6 @@ void streamableHttpTest() {
9494

9595
mcpClient.ping();
9696

97-
System.out.println("mcpClient = " + mcpClient.getServerInfo());
98-
9997
ListToolsResult toolsResult = mcpClient.listTools();
10098

10199
assertThat(toolsResult).isNotNull();

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,6 @@ void streamableHttpTest() {
9595

9696
mcpClient.ping();
9797

98-
System.out.println("mcpClient = " + mcpClient.getServerInfo());
99-
10098
ListToolsResult toolsResult = mcpClient.listTools();
10199

102100
assertThat(toolsResult).isNotNull();

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,6 @@ void streamableHttpTest() {
8181

8282
mcpClient.ping();
8383

84-
System.out.println("mcpClient = " + mcpClient.getServerInfo());
85-
8684
ListToolsResult toolsResult = mcpClient.listTools();
8785

8886
assertThat(toolsResult).isNotNull();

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,8 +82,6 @@ void streamableHttpTest() {
8282

8383
mcpClient.ping();
8484

85-
System.out.println("mcpClient = " + mcpClient.getServerInfo());
86-
8785
ListToolsResult toolsResult = mcpClient.listTools();
8886

8987
assertThat(toolsResult).isNotNull();

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

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,23 +24,30 @@
2424
import org.springaicommunity.mcp.annotation.McpResource;
2525
import org.springaicommunity.mcp.annotation.McpTool;
2626

27+
import org.springframework.ai.mcp.annotation.spring.scan.AbstractAnnotatedMethodBeanFactoryInitializationAotProcessor;
2728
import org.springframework.ai.mcp.annotation.spring.scan.AbstractAnnotatedMethodBeanPostProcessor;
2829
import org.springframework.ai.mcp.annotation.spring.scan.AbstractMcpAnnotatedBeans;
30+
import org.springframework.aot.hint.MemberCategory;
31+
import org.springframework.aot.hint.RuntimeHints;
32+
import org.springframework.aot.hint.RuntimeHintsRegistrar;
2933
import org.springframework.boot.autoconfigure.AutoConfiguration;
3034
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
3135
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
3236
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3337
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3438
import org.springframework.context.annotation.Bean;
39+
import org.springframework.context.annotation.ImportRuntimeHints;
3540

3641
/**
3742
* @author Christian Tzolov
43+
* @author Josh Long
3844
*/
3945
@AutoConfiguration
4046
@ConditionalOnClass(McpTool.class)
4147
@ConditionalOnProperty(prefix = McpServerAnnotationScannerProperties.CONFIG_PREFIX, name = "enabled",
4248
havingValue = "true", matchIfMissing = true)
4349
@EnableConfigurationProperties(McpServerAnnotationScannerProperties.class)
50+
@ImportRuntimeHints(McpServerAnnotationScannerAutoConfiguration.AnnotationHints.class)
4451
public class McpServerAnnotationScannerAutoConfiguration {
4552

4653
private static final Set<Class<? extends Annotation>> SERVER_MCP_ANNOTATIONS = Set.of(McpTool.class,
@@ -54,15 +61,30 @@ public ServerMcpAnnotatedBeans serverAnnotatedBeanRegistry() {
5461

5562
@Bean
5663
@ConditionalOnMissingBean
57-
public ServerAnnotatedMethodBeanPostProcessor serverAnnotatedMethodBeanPostProcessor(
64+
public static ServerAnnotatedMethodBeanPostProcessor serverAnnotatedMethodBeanPostProcessor(
5865
ServerMcpAnnotatedBeans serverMcpAnnotatedBeans, McpServerAnnotationScannerProperties properties) {
5966
return new ServerAnnotatedMethodBeanPostProcessor(serverMcpAnnotatedBeans, SERVER_MCP_ANNOTATIONS);
6067
}
6168

69+
@Bean
70+
public static ServerAnnotatedBeanFactoryInitializationAotProcessor serverAnnotatedBeanFactoryInitializationAotProcessor() {
71+
return new ServerAnnotatedBeanFactoryInitializationAotProcessor(SERVER_MCP_ANNOTATIONS);
72+
}
73+
6274
public static class ServerMcpAnnotatedBeans extends AbstractMcpAnnotatedBeans {
6375

6476
}
6577

78+
public static class ServerAnnotatedBeanFactoryInitializationAotProcessor
79+
extends AbstractAnnotatedMethodBeanFactoryInitializationAotProcessor {
80+
81+
public ServerAnnotatedBeanFactoryInitializationAotProcessor(
82+
Set<Class<? extends Annotation>> targetAnnotations) {
83+
super(targetAnnotations);
84+
}
85+
86+
}
87+
6688
public static class ServerAnnotatedMethodBeanPostProcessor extends AbstractAnnotatedMethodBeanPostProcessor {
6789

6890
public ServerAnnotatedMethodBeanPostProcessor(ServerMcpAnnotatedBeans serverMcpAnnotatedBeans,
@@ -72,4 +94,13 @@ public ServerAnnotatedMethodBeanPostProcessor(ServerMcpAnnotatedBeans serverMcpA
7294

7395
}
7496

97+
static class AnnotationHints implements RuntimeHintsRegistrar {
98+
99+
@Override
100+
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
101+
SERVER_MCP_ANNOTATIONS.forEach(an -> hints.reflection().registerType(an, MemberCategory.values()));
102+
}
103+
104+
}
105+
75106
}

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,10 @@ static class SyncServerSpecificationConfiguration {
5454
@Bean
5555
public List<McpServerFeatures.SyncResourceSpecification> resourceSpecs(
5656
ServerMcpAnnotatedBeans beansWithMcpMethodAnnotations) {
57-
return SyncMcpAnnotationProviders
57+
58+
List<McpServerFeatures.SyncResourceSpecification> syncResourceSpecifications = SyncMcpAnnotationProviders
5859
.resourceSpecifications(beansWithMcpMethodAnnotations.getBeansByAnnotation(McpResource.class));
60+
return syncResourceSpecifications;
5961
}
6062

6163
@Bean
@@ -75,8 +77,10 @@ public List<McpServerFeatures.SyncCompletionSpecification> completionSpecs(
7577
@Bean
7678
public List<McpServerFeatures.SyncToolSpecification> toolSpecs(
7779
ServerMcpAnnotatedBeans beansWithMcpMethodAnnotations) {
78-
return SyncMcpAnnotationProviders
79-
.toolSpecifications(beansWithMcpMethodAnnotations.getBeansByAnnotation(McpTool.class));
80+
List<Object> beansByAnnotation = beansWithMcpMethodAnnotations.getBeansByAnnotation(McpTool.class);
81+
List<McpServerFeatures.SyncToolSpecification> syncToolSpecifications = SyncMcpAnnotationProviders
82+
.toolSpecifications(beansByAnnotation);
83+
return syncToolSpecifications;
8084
}
8185

8286
}

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,10 @@ public List<McpStatelessServerFeatures.SyncCompletionSpecification> completionSp
7777
@Bean
7878
public List<McpStatelessServerFeatures.SyncToolSpecification> toolSpecs(
7979
ServerMcpAnnotatedBeans beansWithMcpMethodAnnotations) {
80-
return SyncMcpAnnotationProviders
81-
.statelessToolSpecifications(beansWithMcpMethodAnnotations.getBeansByAnnotation(McpTool.class));
80+
List<Object> beansByAnnotation = beansWithMcpMethodAnnotations.getBeansByAnnotation(McpTool.class);
81+
List<McpStatelessServerFeatures.SyncToolSpecification> syncToolSpecifications = SyncMcpAnnotationProviders
82+
.statelessToolSpecifications(beansByAnnotation);
83+
return syncToolSpecifications;
8284
}
8385

8486
}

‎mcp/common/src/main/java/org/springframework/ai/mcp/aot/McpHints.java‎

Lines changed: 3 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -16,14 +16,11 @@
1616

1717
package org.springframework.ai.mcp.aot;
1818

19-
import java.util.ArrayList;
20-
import java.util.Arrays;
21-
import java.util.HashSet;
2219
import java.util.Set;
23-
import java.util.stream.Collectors;
2420

2521
import io.modelcontextprotocol.spec.McpSchema;
2622

23+
import org.springframework.ai.aot.AiRuntimeHints;
2724
import org.springframework.aot.hint.MemberCategory;
2825
import org.springframework.aot.hint.RuntimeHints;
2926
import org.springframework.aot.hint.RuntimeHintsRegistrar;
@@ -65,45 +62,10 @@ public class McpHints implements RuntimeHintsRegistrar {
6562
public void registerHints(RuntimeHints hints, @Nullable ClassLoader classLoader) {
6663
var mcs = MemberCategory.values();
6764

68-
for (var tr : innerClasses(McpSchema.class)) {
65+
Set<TypeReference> typeReferences = AiRuntimeHints.findInnerClassesFor(McpSchema.class);
66+
for (var tr : typeReferences) {
6967
hints.reflection().registerType(tr, mcs);
7068
}
7169
}
7270

73-
/**
74-
* Discovers all inner classes of a given class.
75-
* <p>
76-
* This method recursively finds all nested classes (both declared and inherited) of
77-
* the provided class and converts them to type references.
78-
* @param clazz the class to find inner classes for
79-
* @return a set of type references for all discovered inner classes
80-
*/
81-
private Set<TypeReference> innerClasses(Class<?> clazz) {
82-
var indent = new HashSet<String>();
83-
this.findNestedClasses(clazz, indent);
84-
return indent.stream().map(TypeReference::of).collect(Collectors.toSet());
85-
}
86-
87-
/**
88-
* Recursively finds all nested classes of a given class.
89-
* <p>
90-
* This method:
91-
* <ol>
92-
* <li>Collects both declared and inherited nested classes</li>
93-
* <li>Recursively processes each nested class</li>
94-
* <li>Adds the class names to the provided set</li>
95-
* </ol>
96-
* @param clazz the class to find nested classes for
97-
* @param indent the set to collect class names in
98-
*/
99-
private void findNestedClasses(Class<?> clazz, Set<String> indent) {
100-
var classes = new ArrayList<Class<?>>();
101-
classes.addAll(Arrays.asList(clazz.getDeclaredClasses()));
102-
classes.addAll(Arrays.asList(clazz.getClasses()));
103-
for (var nestedClass : classes) {
104-
this.findNestedClasses(nestedClass, indent);
105-
}
106-
indent.addAll(classes.stream().map(Class::getName).toList());
107-
}
108-
10971
}

‎mcp/mcp-annotations-spring/src/main/java/org/springframework/ai/mcp/annotation/spring/scan/AbstractAnnotatedMethodBeanFactoryInitializationAotProcessor.java‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
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.annotation.spring.scan;
18+
19+
import java.lang.annotation.Annotation;
20+
import java.util.ArrayList;
21+
import java.util.List;
22+
import java.util.Set;
23+
24+
import org.springframework.aot.hint.MemberCategory;
25+
import org.springframework.aot.hint.RuntimeHints;
26+
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotContribution;
27+
import org.springframework.beans.factory.aot.BeanFactoryInitializationAotProcessor;
28+
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
29+
import org.springframework.core.log.LogAccessor;
30+
31+
/**
32+
* @author Josh Long
33+
*/
34+
public class AbstractAnnotatedMethodBeanFactoryInitializationAotProcessor extends AnnotatedMethodDiscovery
35+
implements BeanFactoryInitializationAotProcessor {
36+
37+
private static final LogAccessor logger = new LogAccessor(AbstractAnnotatedMethodBeanPostProcessor.class);
38+
39+
public AbstractAnnotatedMethodBeanFactoryInitializationAotProcessor(
40+
Set<Class<? extends Annotation>> targetAnnotations) {
41+
super(targetAnnotations);
42+
}
43+
44+
@Override
45+
public BeanFactoryInitializationAotContribution processAheadOfTime(ConfigurableListableBeanFactory beanFactory) {
46+
List<Class<?>> types = new ArrayList<>();
47+
for (String beanName : beanFactory.getBeanDefinitionNames()) {
48+
Class<?> beanClass = beanFactory.getType(beanName);
49+
Set<Class<? extends Annotation>> classes = this.scan(beanClass);
50+
if (!classes.isEmpty()) {
51+
types.add(beanClass);
52+
}
53+
}
54+
return (generationContext, beanFactoryInitializationCode) -> {
55+
RuntimeHints runtimeHints = generationContext.getRuntimeHints();
56+
for (Class<?> typeReference : types) {
57+
runtimeHints.reflection().registerType(typeReference, MemberCategory.values());
58+
logger.info("registering " + typeReference.getName() + " for reflection");
59+
}
60+
};
61+
}
62+
63+
}

0 commit comments

Comments
 (0)