Skip to content

Commit b99f26d

Browse files
quaffilayaperumalg
authored andcommitted
Allow application to provide custom RouterFunction
Fix GH-4643 Signed-off-by: Yanming Zhou <[email protected]>
1 parent 0f75982 commit b99f26d

File tree

12 files changed

+151
-8
lines changed

12 files changed

+151
-8
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: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@
2828
import org.springframework.web.reactive.function.server.RouterFunction;
2929

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

3234
class McpServerSseWebFluxAutoConfigurationIT {
3335

@@ -96,4 +98,30 @@ void serverBaseUrlConfiguration() {
9698
.isEqualTo("/test"));
9799
}
98100

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

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/McpServerStreamableHttpWebFluxAutoConfigurationIT.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 McpServerStreamableHttpWebFluxAutoConfigurationIT {
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/McpServerSseWebMvcAutoConfiguration.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -90,7 +90,9 @@ public WebMvcSseServerTransportProvider webMvcSseServerTransportProvider(
9090
}
9191

9292
@Bean
93-
public RouterFunction<ServerResponse> mvcMcpRouterFunction(WebMvcSseServerTransportProvider transportProvider) {
93+
@ConditionalOnMissingBean(name = "webMvcSseServerRouterFunction")
94+
public RouterFunction<ServerResponse> webMvcSseServerRouterFunction(
95+
WebMvcSseServerTransportProvider transportProvider) {
9496
return transportProvider.getRouterFunction();
9597
}
9698

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,11 @@
3232
import org.springframework.context.annotation.Bean;
3333
import org.springframework.context.annotation.Conditional;
3434
import org.springframework.web.servlet.function.RouterFunction;
35+
import org.springframework.web.servlet.function.ServerResponse;
3536

3637
/**
3738
* @author Christian Tzolov
39+
* @author Yanming Zhou
3840
*/
3941
@AutoConfiguration(before = McpServerStatelessAutoConfiguration.class)
4042
@ConditionalOnClass({ McpSchema.class })
@@ -60,7 +62,8 @@ public WebMvcStatelessServerTransport webMvcStatelessServerTransport(
6062
// Router function for stateless http transport used by Spring WebFlux to start an
6163
// HTTP server.
6264
@Bean
63-
public RouterFunction<?> webMvcStatelessServerRouterFunction(
65+
@ConditionalOnMissingBean(name = "webMvcStatelessServerRouterFunction")
66+
public RouterFunction<ServerResponse> webMvcStatelessServerRouterFunction(
6467
WebMvcStatelessServerTransport webMvcStatelessTransport) {
6568
return webMvcStatelessTransport.getRouterFunction();
6669
}

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,11 @@
3333
import org.springframework.context.annotation.Bean;
3434
import org.springframework.context.annotation.Conditional;
3535
import org.springframework.web.servlet.function.RouterFunction;
36+
import org.springframework.web.servlet.function.ServerResponse;
3637

3738
/**
3839
* @author Christian Tzolov
40+
* @author Yanming Zhou
3941
*/
4042
@AutoConfiguration(before = McpServerAutoConfiguration.class)
4143
@ConditionalOnClass({ McpSchema.class })
@@ -62,7 +64,8 @@ public WebMvcStreamableServerTransportProvider webMvcStreamableServerTransportPr
6264
// Router function for streamable http transport used by Spring WebFlux to start an
6365
// HTTP server.
6466
@Bean
65-
public RouterFunction<?> webMvcStreamableServerRouterFunction(
67+
@ConditionalOnMissingBean(name = "webMvcStreamableServerRouterFunction")
68+
public RouterFunction<ServerResponse> webMvcStreamableServerRouterFunction(
6669
WebMvcStreamableServerTransportProvider webMvcProvider) {
6770
return webMvcProvider.getRouterFunction();
6871
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@
3232
import org.springframework.web.servlet.function.RouterFunction;
3333

3434
import static org.assertj.core.api.Assertions.assertThat;
35+
import static org.mockito.Mockito.mock;
36+
import static org.mockito.Mockito.mockingDetails;
3537

3638
class McpServerSseWebMvcAutoConfigurationIT {
3739

@@ -117,4 +119,29 @@ public ConfigurableEnvironment getEnvironment() {
117119
});
118120
}
119121

122+
@Test
123+
void routerFunctionIsCreatedFromProvider() {
124+
this.contextRunner.run(context -> {
125+
assertThat(context).hasSingleBean(RouterFunction.class);
126+
assertThat(context).hasSingleBean(WebMvcSseServerTransportProvider.class);
127+
128+
// Verify that the RouterFunction is created from the provider
129+
WebMvcSseServerTransportProvider serverTransport = context.getBean(WebMvcSseServerTransportProvider.class);
130+
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
131+
assertThat(routerFunction).isNotNull().isEqualTo(serverTransport.getRouterFunction());
132+
});
133+
}
134+
135+
@Test
136+
void routerFunctionIsCustom() {
137+
this.contextRunner
138+
.withBean("webMvcSseServerRouterFunction", RouterFunction.class, () -> mock(RouterFunction.class))
139+
.run(context -> {
140+
assertThat(context).hasSingleBean(RouterFunction.class);
141+
142+
RouterFunction<?> routerFunction = context.getBean(RouterFunction.class);
143+
assertThat(mockingDetails(routerFunction).isMock()).isTrue();
144+
});
145+
}
146+
120147
}

0 commit comments

Comments
 (0)