Skip to content

Commit 698996e

Browse files
quaffquaff
committed
Allow application to provide custom RouterFunction
Fix GH-4643 Signed-off-by: Yanming Zhou <[email protected]>
1 parent 940bcf3 commit 698996e

File tree

10 files changed

+118
-5
lines changed

10 files changed

+118
-5
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,8 @@ public WebFluxSseServerTransportProvider webFluxTransport(ObjectProvider<ObjectM
9999
// Router function for SSE transport used by Spring WebFlux to start an HTTP
100100
// server.
101101
@Bean
102-
public RouterFunction<?> webfluxMcpRouterFunction(WebFluxSseServerTransportProvider webFluxProvider) {
102+
@ConditionalOnMissingBean(name = "webfluxSseServerRouterFunction")
103+
public RouterFunction<?> webfluxSseServerRouterFunction(WebFluxSseServerTransportProvider webFluxProvider) {
103104
return webFluxProvider.getRouterFunction();
104105
}
105106

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
/**
3737
* @author Christian Tzolov
38+
* @author Yanming Zhou
3839
*/
3940
@AutoConfiguration(before = McpServerStatelessAutoConfiguration.class)
4041
@ConditionalOnClass({ McpSchema.class })
@@ -60,6 +61,7 @@ public WebFluxStatelessServerTransport webFluxStatelessServerTransport(
6061
// Router function for stateless http transport used by Spring WebFlux to start an
6162
// HTTP server.
6263
@Bean
64+
@ConditionalOnMissingBean(name = "webFluxStatelessServerRouterFunction")
6365
public RouterFunction<?> webFluxStatelessServerRouterFunction(
6466
WebFluxStatelessServerTransport webFluxStatelessTransport) {
6567
return webFluxStatelessTransport.getRouterFunction();

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636

3737
/**
3838
* @author Christian Tzolov
39+
* @author Yanming Zhou
3940
*/
4041
@AutoConfiguration(before = McpServerAutoConfiguration.class)
4142
@ConditionalOnClass({ McpSchema.class })
@@ -62,6 +63,7 @@ public WebFluxStreamableServerTransportProvider webFluxStreamableServerTransport
6263
// Router function for streamable http transport used by Spring WebFlux to start an
6364
// HTTP server.
6465
@Bean
66+
@ConditionalOnMissingBean(name = "webFluxStreamableServerRouterFunction")
6567
public RouterFunction<?> webFluxStreamableServerRouterFunction(
6668
WebFluxStreamableServerTransportProvider webFluxProvider) {
6769
return webFluxProvider.getRouterFunction();

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.fasterxml.jackson.databind.ObjectMapper;
2020
import io.modelcontextprotocol.server.McpSyncServer;
2121
import io.modelcontextprotocol.server.transport.WebFluxSseServerTransportProvider;
22+
import io.modelcontextprotocol.server.transport.WebFluxStatelessServerTransport;
2223
import org.junit.jupiter.api.Test;
2324

2425
import org.springframework.ai.mcp.server.common.autoconfigure.McpServerAutoConfiguration;
@@ -28,6 +29,8 @@
2829
import org.springframework.web.reactive.function.server.RouterFunction;
2930

3031
import static org.assertj.core.api.Assertions.assertThat;
32+
import static org.mockito.Mockito.mock;
33+
import static org.mockito.Mockito.mockingDetails;
3134

3235
class McpServerSseWebFluxAutoConfigurationIT {
3336

@@ -96,4 +99,30 @@ void serverBaseUrlConfiguration() {
9699
.isEqualTo("/test"));
97100
}
98101

102+
@Test
103+
void routerFunctionIsCreatedFromProvider() {
104+
this.contextRunner.run(context -> {
105+
assertThat(context).hasSingleBean(RouterFunction.class);
106+
assertThat(context).hasSingleBean(WebFluxSseServerTransportProvider.class);
107+
108+
// Verify that the RouterFunction is created from the provider
109+
WebFluxSseServerTransportProvider serverTransport = context
110+
.getBean(WebFluxSseServerTransportProvider.class);
111+
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
112+
assertThat(routerFunction).isNotNull().isEqualTo(serverTransport.getRouterFunction());
113+
});
114+
}
115+
116+
@Test
117+
void routerFunctionIsCustom() {
118+
this.contextRunner
119+
.withBean("webfluxSseServerRouterFunction", RouterFunction.class, () -> mock(RouterFunction.class))
120+
.run(context -> {
121+
assertThat(context).hasSingleBean(RouterFunction.class);
122+
123+
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
124+
assertThat(mockingDetails(routerFunction).isMock()).isTrue();
125+
});
126+
}
127+
99128
}

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

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,13 @@
2323

2424
import org.springframework.boot.autoconfigure.AutoConfigurations;
2525
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
26+
import org.springframework.context.annotation.Bean;
27+
import org.springframework.context.annotation.Configuration;
2628
import org.springframework.web.reactive.function.server.RouterFunction;
2729

2830
import static org.assertj.core.api.Assertions.assertThat;
31+
import static org.mockito.Mockito.mock;
32+
import static org.mockito.Mockito.mockingDetails;
2933

3034
class McpServerStatelessWebFluxAutoConfigurationIT {
3135

@@ -135,11 +139,24 @@ void routerFunctionIsCreatedFromProvider() {
135139
assertThat(context).hasSingleBean(WebFluxStatelessServerTransport.class);
136140

137141
// Verify that the RouterFunction is created from the provider
142+
WebFluxStatelessServerTransport serverTransport = context.getBean(WebFluxStatelessServerTransport.class);
138143
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
139-
assertThat(routerFunction).isNotNull();
144+
assertThat(routerFunction).isNotNull().isEqualTo(serverTransport.getRouterFunction());
140145
});
141146
}
142147

148+
@Test
149+
void routerFunctionIsCustom() {
150+
this.contextRunner
151+
.withBean("webFluxStatelessServerRouterFunction", RouterFunction.class, () -> mock(RouterFunction.class))
152+
.run(context -> {
153+
assertThat(context).hasSingleBean(RouterFunction.class);
154+
155+
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
156+
assertThat(mockingDetails(routerFunction).isMock()).isTrue();
157+
});
158+
}
159+
143160
@Test
144161
void allPropertiesConfiguration() {
145162
this.contextRunner
@@ -172,4 +189,15 @@ void enabledPropertyExplicitlyTrue() {
172189
});
173190
}
174191

192+
@Configuration
193+
private static class CustomRouterFunctionConfig {
194+
195+
@Bean
196+
public RouterFunction<?> webFluxStatelessServerRouterFunction(
197+
WebFluxStatelessServerTransport webFluxStatelessTransport) {
198+
return mock(RouterFunction.class);
199+
}
200+
201+
}
202+
175203
}

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import org.springframework.web.reactive.function.server.RouterFunction;
2727

2828
import static org.assertj.core.api.Assertions.assertThat;
29+
import static org.mockito.Mockito.mock;
30+
import static org.mockito.Mockito.mockingDetails;
2931

3032
class McpServerStreamableWebFluxAutoConfigurationIT {
3133

@@ -137,11 +139,25 @@ void routerFunctionIsCreatedFromProvider() {
137139
assertThat(context).hasSingleBean(WebFluxStreamableServerTransportProvider.class);
138140

139141
// Verify that the RouterFunction is created from the provider
142+
WebFluxStreamableServerTransportProvider serverTransport = context
143+
.getBean(WebFluxStreamableServerTransportProvider.class);
140144
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
141-
assertThat(routerFunction).isNotNull();
145+
assertThat(routerFunction).isNotNull().isEqualTo(serverTransport.getRouterFunction());
142146
});
143147
}
144148

149+
@Test
150+
void routerFunctionIsCustom() {
151+
this.contextRunner
152+
.withBean("webFluxStreamableServerRouterFunction", RouterFunction.class, () -> mock(RouterFunction.class))
153+
.run(context -> {
154+
assertThat(context).hasSingleBean(RouterFunction.class);
155+
156+
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
157+
assertThat(mockingDetails(routerFunction).isMock()).isTrue();
158+
});
159+
}
160+
145161
@Test
146162
void allPropertiesConfiguration() {
147163
this.contextRunner

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
/**
3737
* @author Christian Tzolov
38+
* @author Yanming Zhou
3839
*/
3940
@AutoConfiguration(before = McpServerStatelessAutoConfiguration.class)
4041
@ConditionalOnClass({ McpSchema.class })
@@ -60,6 +61,7 @@ public WebMvcStatelessServerTransport webMvcStatelessServerTransport(
6061
// Router function for stateless http transport used by Spring WebFlux to start an
6162
// HTTP server.
6263
@Bean
64+
@ConditionalOnMissingBean(name = "webMvcStatelessServerRouterFunction")
6365
public RouterFunction<?> webMvcStatelessServerRouterFunction(
6466
WebMvcStatelessServerTransport webMvcStatelessTransport) {
6567
return webMvcStatelessTransport.getRouterFunction();

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636

3737
/**
3838
* @author Christian Tzolov
39+
* @author Yanming Zhou
3940
*/
4041
@AutoConfiguration(before = McpServerAutoConfiguration.class)
4142
@ConditionalOnClass({ McpSchema.class })
@@ -62,6 +63,7 @@ public WebMvcStreamableServerTransportProvider webMvcStreamableServerTransportPr
6263
// Router function for streamable http transport used by Spring WebFlux to start an
6364
// HTTP server.
6465
@Bean
66+
@ConditionalOnMissingBean(name = "webMvcStreamableServerRouterFunction")
6567
public RouterFunction<?> webMvcStreamableServerRouterFunction(
6668
WebMvcStreamableServerTransportProvider webMvcProvider) {
6769
return webMvcProvider.getRouterFunction();

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

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import org.springframework.web.servlet.function.RouterFunction;
2727

2828
import static org.assertj.core.api.Assertions.assertThat;
29+
import static org.mockito.Mockito.mock;
30+
import static org.mockito.Mockito.mockingDetails;
2931

3032
class McpServerStatelessWebMvcAutoConfigurationIT {
3133

@@ -135,11 +137,24 @@ void routerFunctionIsCreatedFromProvider() {
135137
assertThat(context).hasSingleBean(WebMvcStatelessServerTransport.class);
136138

137139
// Verify that the RouterFunction is created from the provider
140+
WebMvcStatelessServerTransport serverTransport = context.getBean(WebMvcStatelessServerTransport.class);
138141
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
139-
assertThat(routerFunction).isNotNull();
142+
assertThat(routerFunction).isNotNull().isEqualTo(serverTransport.getRouterFunction());
140143
});
141144
}
142145

146+
@Test
147+
void routerFunctionIsCustom() {
148+
this.contextRunner
149+
.withBean("webMvcStatelessServerRouterFunction", RouterFunction.class, () -> mock(RouterFunction.class))
150+
.run(context -> {
151+
assertThat(context).hasSingleBean(RouterFunction.class);
152+
153+
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
154+
assertThat(mockingDetails(routerFunction).isMock()).isTrue();
155+
});
156+
}
157+
143158
@Test
144159
void allPropertiesConfiguration() {
145160
this.contextRunner

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,8 @@
2626
import org.springframework.web.servlet.function.RouterFunction;
2727

2828
import static org.assertj.core.api.Assertions.assertThat;
29+
import static org.mockito.Mockito.mock;
30+
import static org.mockito.Mockito.mockingDetails;
2931

3032
class McpServerStreamableWebMvcAutoConfigurationIT {
3133

@@ -137,11 +139,25 @@ void routerFunctionIsCreatedFromProvider() {
137139
assertThat(context).hasSingleBean(WebMvcStreamableServerTransportProvider.class);
138140

139141
// Verify that the RouterFunction is created from the provider
142+
WebMvcStreamableServerTransportProvider serverTransportProvider = context
143+
.getBean(WebMvcStreamableServerTransportProvider.class);
140144
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
141-
assertThat(routerFunction).isNotNull();
145+
assertThat(routerFunction).isNotNull().isEqualTo(serverTransportProvider.getRouterFunction());
142146
});
143147
}
144148

149+
@Test
150+
void routerFunctionIsCustom() {
151+
this.contextRunner
152+
.withBean("webMvcStreamableServerRouterFunction", RouterFunction.class, () -> mock(RouterFunction.class))
153+
.run(context -> {
154+
assertThat(context).hasSingleBean(RouterFunction.class);
155+
156+
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
157+
assertThat(mockingDetails(routerFunction).isMock()).isTrue();
158+
});
159+
}
160+
145161
@Test
146162
void allPropertiesConfiguration() {
147163
this.contextRunner

0 commit comments

Comments
 (0)