|
1 | 1 | /* |
2 | | -* Copyright 2024 - 2024 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 | | -*/ |
| 2 | + * Copyright 2024 - 2024 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 | 16 | package org.springframework.ai.mcp.server; |
17 | 17 |
|
18 | 18 | import java.time.Duration; |
|
46 | 46 | import org.springframework.web.reactive.function.server.RouterFunctions; |
47 | 47 |
|
48 | 48 | import static org.assertj.core.api.Assertions.assertThat; |
| 49 | +import static org.assertj.core.api.Assertions.assertThatThrownBy; |
49 | 50 | import static org.awaitility.Awaitility.await; |
50 | 51 |
|
51 | 52 | public class SseAsyncIntegrationTests { |
@@ -165,20 +166,13 @@ void testCreateMessageSuccess() throws InterruptedException { |
165 | 166 | // --------------------------------------- |
166 | 167 | @Test |
167 | 168 | void testRootsSuccess() { |
168 | | - |
169 | 169 | List<Root> roots = List.of(new Root("uri1://", "root1"), new Root("uri2://", "root2")); |
170 | 170 |
|
171 | 171 | AtomicReference<List<Root>> rootsRef = new AtomicReference<>(); |
172 | 172 | var mcpServer = McpServer.using(mcpServerTransport) |
173 | 173 | .rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate)) |
174 | 174 | .sync(); |
175 | 175 |
|
176 | | - // HttpHandler httpHandler = |
177 | | - // RouterFunctions.toHttpHandler(mcpServerTransport.getRouterFunction()); |
178 | | - // ReactorHttpHandlerAdapter adapter = new ReactorHttpHandlerAdapter(httpHandler); |
179 | | - // HttpServer httpServer = HttpServer.create().port(8080).handle(adapter); |
180 | | - // DisposableServer d = httpServer.bindNow(); |
181 | | - |
182 | 176 | var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) |
183 | 177 | .roots(roots) |
184 | 178 | .sync(); |
@@ -212,8 +206,112 @@ void testRootsSuccess() { |
212 | 206 | }); |
213 | 207 |
|
214 | 208 | mcpClient.close(); |
| 209 | + mcpServer.close(); |
| 210 | + } |
| 211 | + |
| 212 | + @Test |
| 213 | + void testRootsWithoutCapability() { |
| 214 | + var mcpServer = McpServer.using(mcpServerTransport).rootsChangeConsumer(rootsUpdate -> { |
| 215 | + }).sync(); |
| 216 | + |
| 217 | + // Create client without roots capability |
| 218 | + var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().build()) // No |
| 219 | + // roots |
| 220 | + // capability |
| 221 | + .sync(); |
| 222 | + |
| 223 | + InitializeResult initResult = mcpClient.initialize(); |
| 224 | + assertThat(initResult).isNotNull(); |
| 225 | + |
| 226 | + // Attempt to list roots should fail |
| 227 | + assertThatThrownBy(() -> mcpServer.listRoots().roots()).isInstanceOf(McpError.class) |
| 228 | + .hasMessage("Roots not supported"); |
| 229 | + |
| 230 | + mcpClient.close(); |
| 231 | + mcpServer.close(); |
| 232 | + } |
| 233 | + |
| 234 | + @Test |
| 235 | + void testRootsWithEmptyRootsList() { |
| 236 | + AtomicReference<List<Root>> rootsRef = new AtomicReference<>(); |
| 237 | + var mcpServer = McpServer.using(mcpServerTransport) |
| 238 | + .rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate)) |
| 239 | + .sync(); |
| 240 | + |
| 241 | + var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) |
| 242 | + .roots(List.of()) // Empty roots list |
| 243 | + .sync(); |
| 244 | + |
| 245 | + InitializeResult initResult = mcpClient.initialize(); |
| 246 | + assertThat(initResult).isNotNull(); |
| 247 | + |
| 248 | + mcpClient.rootsListChangedNotification(); |
| 249 | + |
| 250 | + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { |
| 251 | + assertThat(rootsRef.get()).isEmpty(); |
| 252 | + }); |
| 253 | + |
| 254 | + mcpClient.close(); |
| 255 | + mcpServer.close(); |
| 256 | + } |
| 257 | + |
| 258 | + @Test |
| 259 | + void testRootsWithMultipleConsumers() { |
| 260 | + List<Root> roots = List.of(new Root("uri1://", "root1")); |
| 261 | + |
| 262 | + AtomicReference<List<Root>> rootsRef1 = new AtomicReference<>(); |
| 263 | + AtomicReference<List<Root>> rootsRef2 = new AtomicReference<>(); |
| 264 | + |
| 265 | + var mcpServer = McpServer.using(mcpServerTransport) |
| 266 | + .rootsChangeConsumer(rootsUpdate -> rootsRef1.set(rootsUpdate)) |
| 267 | + .rootsChangeConsumer(rootsUpdate -> rootsRef2.set(rootsUpdate)) |
| 268 | + .sync(); |
| 269 | + |
| 270 | + var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) |
| 271 | + .roots(roots) |
| 272 | + .sync(); |
| 273 | + |
| 274 | + InitializeResult initResult = mcpClient.initialize(); |
| 275 | + assertThat(initResult).isNotNull(); |
| 276 | + |
| 277 | + mcpClient.rootsListChangedNotification(); |
| 278 | + |
| 279 | + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { |
| 280 | + assertThat(rootsRef1.get()).containsAll(roots); |
| 281 | + assertThat(rootsRef2.get()).containsAll(roots); |
| 282 | + }); |
| 283 | + |
| 284 | + mcpClient.close(); |
| 285 | + mcpServer.close(); |
| 286 | + } |
| 287 | + |
| 288 | + @Test |
| 289 | + void testRootsServerCloseWithActiveSubscription() { |
| 290 | + List<Root> roots = List.of(new Root("uri1://", "root1")); |
215 | 291 |
|
| 292 | + AtomicReference<List<Root>> rootsRef = new AtomicReference<>(); |
| 293 | + var mcpServer = McpServer.using(mcpServerTransport) |
| 294 | + .rootsChangeConsumer(rootsUpdate -> rootsRef.set(rootsUpdate)) |
| 295 | + .sync(); |
| 296 | + |
| 297 | + var mcpClient = clientBuilder.capabilities(ClientCapabilities.builder().roots(true).build()) |
| 298 | + .roots(roots) |
| 299 | + .sync(); |
| 300 | + |
| 301 | + InitializeResult initResult = mcpClient.initialize(); |
| 302 | + assertThat(initResult).isNotNull(); |
| 303 | + |
| 304 | + mcpClient.rootsListChangedNotification(); |
| 305 | + |
| 306 | + await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> { |
| 307 | + assertThat(rootsRef.get()).containsAll(roots); |
| 308 | + }); |
| 309 | + |
| 310 | + // Close server while subscription is active |
216 | 311 | mcpServer.close(); |
| 312 | + |
| 313 | + // Verify client can handle server closure gracefully |
| 314 | + mcpClient.close(); |
217 | 315 | } |
218 | 316 |
|
219 | 317 | } |
0 commit comments