|
4 | 4 | package io.modelcontextprotocol;
|
5 | 5 |
|
6 | 6 | import java.time.Duration;
|
7 |
| -import java.util.ArrayList; |
8 | 7 | import java.util.List;
|
9 | 8 | import java.util.Map;
|
10 | 9 | import java.util.concurrent.ConcurrentHashMap;
|
|
28 | 27 | import io.modelcontextprotocol.spec.McpError;
|
29 | 28 | import io.modelcontextprotocol.spec.McpSchema;
|
30 | 29 | import io.modelcontextprotocol.spec.McpSchema.*;
|
31 |
| -import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities.CompletionCapabilities; |
32 | 30 | import org.junit.jupiter.api.AfterEach;
|
33 | 31 | import org.junit.jupiter.api.BeforeEach;
|
34 | 32 | import org.junit.jupiter.params.ParameterizedTest;
|
35 | 33 | import org.junit.jupiter.params.provider.ValueSource;
|
| 34 | +import reactor.core.publisher.Mono; |
36 | 35 | import reactor.netty.DisposableServer;
|
37 | 36 | import reactor.netty.http.server.HttpServer;
|
38 | 37 |
|
|
41 | 40 | import org.springframework.web.client.RestClient;
|
42 | 41 | import org.springframework.web.reactive.function.client.WebClient;
|
43 | 42 | import org.springframework.web.reactive.function.server.RouterFunctions;
|
| 43 | +import reactor.test.StepVerifier; |
44 | 44 |
|
45 | 45 | import static org.assertj.core.api.Assertions.assertThat;
|
46 | 46 | import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
|
@@ -331,6 +331,226 @@ void testCreateMessageWithRequestTimeoutFail(String clientType) throws Interrupt
|
331 | 331 | mcpServer.closeGracefully().block();
|
332 | 332 | }
|
333 | 333 |
|
| 334 | + // --------------------------------------- |
| 335 | + // Elicitation Tests |
| 336 | + // --------------------------------------- |
| 337 | + @ParameterizedTest(name = "{0} : {displayName} ") |
| 338 | + @ValueSource(strings = { "httpclient", "webflux" }) |
| 339 | + void testCreateElicitationWithoutElicitationCapabilities(String clientType) { |
| 340 | + |
| 341 | + var clientBuilder = clientBuilders.get(clientType); |
| 342 | + |
| 343 | + McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification( |
| 344 | + new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> { |
| 345 | + |
| 346 | + exchange.createElicitation(mock(ElicitRequest.class)).block(); |
| 347 | + |
| 348 | + return Mono.just(mock(CallToolResult.class)); |
| 349 | + }); |
| 350 | + |
| 351 | + var server = McpServer.async(mcpServerTransportProvider).serverInfo("test-server", "1.0.0").tools(tool).build(); |
| 352 | + |
| 353 | + try ( |
| 354 | + // Create client without elicitation capabilities |
| 355 | + var client = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")).build()) { |
| 356 | + |
| 357 | + assertThat(client.initialize()).isNotNull(); |
| 358 | + |
| 359 | + try { |
| 360 | + client.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); |
| 361 | + } |
| 362 | + catch (McpError e) { |
| 363 | + assertThat(e).isInstanceOf(McpError.class) |
| 364 | + .hasMessage("Client must be configured with elicitation capabilities"); |
| 365 | + } |
| 366 | + } |
| 367 | + server.closeGracefully().block(); |
| 368 | + } |
| 369 | + |
| 370 | + @ParameterizedTest(name = "{0} : {displayName} ") |
| 371 | + @ValueSource(strings = { "httpclient", "webflux" }) |
| 372 | + void testCreateElicitationSuccess(String clientType) { |
| 373 | + |
| 374 | + var clientBuilder = clientBuilders.get(clientType); |
| 375 | + |
| 376 | + Function<ElicitRequest, ElicitResult> elicitationHandler = request -> { |
| 377 | + assertThat(request.message()).isNotEmpty(); |
| 378 | + assertThat(request.requestedSchema()).isNotNull(); |
| 379 | + |
| 380 | + return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); |
| 381 | + }; |
| 382 | + |
| 383 | + CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), |
| 384 | + null); |
| 385 | + |
| 386 | + McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification( |
| 387 | + new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> { |
| 388 | + |
| 389 | + var elicitationRequest = ElicitRequest.builder() |
| 390 | + .message("Test message") |
| 391 | + .requestedSchema( |
| 392 | + Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) |
| 393 | + .build(); |
| 394 | + |
| 395 | + StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { |
| 396 | + assertThat(result).isNotNull(); |
| 397 | + assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); |
| 398 | + assertThat(result.content().get("message")).isEqualTo("Test message"); |
| 399 | + }).verifyComplete(); |
| 400 | + |
| 401 | + return Mono.just(callResponse); |
| 402 | + }); |
| 403 | + |
| 404 | + var mcpServer = McpServer.async(mcpServerTransportProvider) |
| 405 | + .serverInfo("test-server", "1.0.0") |
| 406 | + .tools(tool) |
| 407 | + .build(); |
| 408 | + |
| 409 | + try (var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) |
| 410 | + .capabilities(ClientCapabilities.builder().elicitation().build()) |
| 411 | + .elicitation(elicitationHandler) |
| 412 | + .build()) { |
| 413 | + |
| 414 | + InitializeResult initResult = mcpClient.initialize(); |
| 415 | + assertThat(initResult).isNotNull(); |
| 416 | + |
| 417 | + CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); |
| 418 | + |
| 419 | + assertThat(response).isNotNull(); |
| 420 | + assertThat(response).isEqualTo(callResponse); |
| 421 | + } |
| 422 | + mcpServer.closeGracefully().block(); |
| 423 | + } |
| 424 | + |
| 425 | + @ParameterizedTest(name = "{0} : {displayName} ") |
| 426 | + @ValueSource(strings = { "httpclient", "webflux" }) |
| 427 | + void testCreateElicitationWithRequestTimeoutSuccess(String clientType) { |
| 428 | + |
| 429 | + // Client |
| 430 | + var clientBuilder = clientBuilders.get(clientType); |
| 431 | + |
| 432 | + Function<ElicitRequest, ElicitResult> elicitationHandler = request -> { |
| 433 | + assertThat(request.message()).isNotEmpty(); |
| 434 | + assertThat(request.requestedSchema()).isNotNull(); |
| 435 | + try { |
| 436 | + TimeUnit.SECONDS.sleep(2); |
| 437 | + } |
| 438 | + catch (InterruptedException e) { |
| 439 | + throw new RuntimeException(e); |
| 440 | + } |
| 441 | + return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); |
| 442 | + }; |
| 443 | + |
| 444 | + var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) |
| 445 | + .capabilities(ClientCapabilities.builder().elicitation().build()) |
| 446 | + .elicitation(elicitationHandler) |
| 447 | + .build(); |
| 448 | + |
| 449 | + // Server |
| 450 | + |
| 451 | + CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), |
| 452 | + null); |
| 453 | + |
| 454 | + McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification( |
| 455 | + new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> { |
| 456 | + |
| 457 | + var elicitationRequest = ElicitRequest.builder() |
| 458 | + .message("Test message") |
| 459 | + .requestedSchema( |
| 460 | + Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) |
| 461 | + .build(); |
| 462 | + |
| 463 | + StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { |
| 464 | + assertThat(result).isNotNull(); |
| 465 | + assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); |
| 466 | + assertThat(result.content().get("message")).isEqualTo("Test message"); |
| 467 | + }).verifyComplete(); |
| 468 | + |
| 469 | + return Mono.just(callResponse); |
| 470 | + }); |
| 471 | + |
| 472 | + var mcpServer = McpServer.async(mcpServerTransportProvider) |
| 473 | + .serverInfo("test-server", "1.0.0") |
| 474 | + .requestTimeout(Duration.ofSeconds(3)) |
| 475 | + .tools(tool) |
| 476 | + .build(); |
| 477 | + |
| 478 | + InitializeResult initResult = mcpClient.initialize(); |
| 479 | + assertThat(initResult).isNotNull(); |
| 480 | + |
| 481 | + CallToolResult response = mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); |
| 482 | + |
| 483 | + assertThat(response).isNotNull(); |
| 484 | + assertThat(response).isEqualTo(callResponse); |
| 485 | + |
| 486 | + mcpClient.closeGracefully(); |
| 487 | + mcpServer.closeGracefully().block(); |
| 488 | + } |
| 489 | + |
| 490 | + @ParameterizedTest(name = "{0} : {displayName} ") |
| 491 | + @ValueSource(strings = { "httpclient", "webflux" }) |
| 492 | + void testCreateElicitationWithRequestTimeoutFail(String clientType) { |
| 493 | + |
| 494 | + // Client |
| 495 | + var clientBuilder = clientBuilders.get(clientType); |
| 496 | + |
| 497 | + Function<ElicitRequest, ElicitResult> elicitationHandler = request -> { |
| 498 | + assertThat(request.message()).isNotEmpty(); |
| 499 | + assertThat(request.requestedSchema()).isNotNull(); |
| 500 | + try { |
| 501 | + TimeUnit.SECONDS.sleep(2); |
| 502 | + } |
| 503 | + catch (InterruptedException e) { |
| 504 | + throw new RuntimeException(e); |
| 505 | + } |
| 506 | + return new ElicitResult(ElicitResult.Action.ACCEPT, Map.of("message", request.message())); |
| 507 | + }; |
| 508 | + |
| 509 | + var mcpClient = clientBuilder.clientInfo(new McpSchema.Implementation("Sample client", "0.0.0")) |
| 510 | + .capabilities(ClientCapabilities.builder().elicitation().build()) |
| 511 | + .elicitation(elicitationHandler) |
| 512 | + .build(); |
| 513 | + |
| 514 | + // Server |
| 515 | + |
| 516 | + CallToolResult callResponse = new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("CALL RESPONSE")), |
| 517 | + null); |
| 518 | + |
| 519 | + McpServerFeatures.AsyncToolSpecification tool = new McpServerFeatures.AsyncToolSpecification( |
| 520 | + new McpSchema.Tool("tool1", "tool1 description", emptyJsonSchema), (exchange, request) -> { |
| 521 | + |
| 522 | + var elicitationRequest = ElicitRequest.builder() |
| 523 | + .message("Test message") |
| 524 | + .requestedSchema( |
| 525 | + Map.of("type", "object", "properties", Map.of("message", Map.of("type", "string")))) |
| 526 | + .build(); |
| 527 | + |
| 528 | + StepVerifier.create(exchange.createElicitation(elicitationRequest)).consumeNextWith(result -> { |
| 529 | + assertThat(result).isNotNull(); |
| 530 | + assertThat(result.action()).isEqualTo(ElicitResult.Action.ACCEPT); |
| 531 | + assertThat(result.content().get("message")).isEqualTo("Test message"); |
| 532 | + }).verifyComplete(); |
| 533 | + |
| 534 | + return Mono.just(callResponse); |
| 535 | + }); |
| 536 | + |
| 537 | + var mcpServer = McpServer.async(mcpServerTransportProvider) |
| 538 | + .serverInfo("test-server", "1.0.0") |
| 539 | + .requestTimeout(Duration.ofSeconds(1)) |
| 540 | + .tools(tool) |
| 541 | + .build(); |
| 542 | + |
| 543 | + InitializeResult initResult = mcpClient.initialize(); |
| 544 | + assertThat(initResult).isNotNull(); |
| 545 | + |
| 546 | + assertThatExceptionOfType(McpError.class).isThrownBy(() -> { |
| 547 | + mcpClient.callTool(new McpSchema.CallToolRequest("tool1", Map.of())); |
| 548 | + }).withMessageContaining("within 1000ms"); |
| 549 | + |
| 550 | + mcpClient.closeGracefully(); |
| 551 | + mcpServer.closeGracefully().block(); |
| 552 | + } |
| 553 | + |
334 | 554 | // ---------------------------------------
|
335 | 555 | // Roots Tests
|
336 | 556 | // ---------------------------------------
|
|
0 commit comments