feat(client): add client tool output schema validation and caching (#302)
- Add JSON schema validation for tool call results against output schemas - Implement automatic tool output schema caching during initialization - Add `enableCallToolSchemaCaching` configuration option to enable/disable schema caching - Add `JsonSchemaValidator` integration to McpClient builder APIs - Introduce post-initialization hook mechanism for performing operations after successful client initialization - Cache tool output schemas during `listTools` operations when caching is enabled - Validate tool results against cached schemas in `callTool` operations - Return error CallToolResult when validation fails - Add test coverage - Convert validateToolResult from Mono to synchronous method - Throw IllegalArgumentException on validation errors Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com> Co-authored-by: Anurag Pant <pantanurag555@github> Co-authored-by: Christian Tzolov <christian.tzolov@broadcom.com> Co-authored-by: Daniel Garnier-Moiroux <git@garnier.wf> Co-authored-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
Anurag Pant
Christian Tzolov
parent
3f7578b3ce
commit
8e11fa82ae
5 changed files with 708 additions and 131 deletions
|
|
@ -4,15 +4,14 @@
|
|||
|
||||
package io.modelcontextprotocol.client;
|
||||
|
||||
import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
import org.junit.jupiter.api.AfterAll;
|
||||
import org.junit.jupiter.api.BeforeAll;
|
||||
import org.junit.jupiter.api.Timeout;
|
||||
import org.testcontainers.containers.GenericContainer;
|
||||
import org.testcontainers.containers.wait.strategy.Wait;
|
||||
|
||||
import io.modelcontextprotocol.client.transport.HttpClientStreamableHttpTransport;
|
||||
import io.modelcontextprotocol.spec.McpClientTransport;
|
||||
|
||||
@Timeout(15)
|
||||
public class HttpClientStreamableHttpAsyncClientTests extends AbstractMcpAsyncClientTests {
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,280 @@
|
|||
/*
|
||||
* Copyright 2024-2024 the original author or authors.
|
||||
*/
|
||||
|
||||
package io.modelcontextprotocol.client;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
|
||||
import io.modelcontextprotocol.spec.McpClientSession;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.MockitoAnnotations;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
import reactor.test.StepVerifier;
|
||||
import reactor.util.context.ContextView;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
/**
|
||||
* Tests for {@link LifecycleInitializer} postInitializationHook functionality.
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
*/
|
||||
class LifecycleInitializerPostInitializationHookTests {
|
||||
|
||||
private static final Duration INITIALIZATION_TIMEOUT = Duration.ofSeconds(5);
|
||||
|
||||
private static final McpSchema.ClientCapabilities CLIENT_CAPABILITIES = McpSchema.ClientCapabilities.builder()
|
||||
.build();
|
||||
|
||||
private static final McpSchema.Implementation CLIENT_INFO = new McpSchema.Implementation("test-client", "1.0.0");
|
||||
|
||||
private static final List<String> PROTOCOL_VERSIONS = List.of("1.0.0", "2.0.0");
|
||||
|
||||
private static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult("2.0.0",
|
||||
McpSchema.ServerCapabilities.builder().build(), new McpSchema.Implementation("test-server", "1.0.0"),
|
||||
"Test instructions");
|
||||
|
||||
@Mock
|
||||
private McpClientSession mockClientSession;
|
||||
|
||||
@Mock
|
||||
private Function<ContextView, McpClientSession> mockSessionSupplier;
|
||||
|
||||
@Mock
|
||||
private Function<Initialization, Mono<Void>> mockPostInitializationHook;
|
||||
|
||||
private LifecycleInitializer initializer;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
MockitoAnnotations.openMocks(this);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.empty());
|
||||
when(mockSessionSupplier.apply(any(ContextView.class))).thenReturn(mockClientSession);
|
||||
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.just(MOCK_INIT_RESULT));
|
||||
when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
|
||||
.thenReturn(Mono.empty());
|
||||
when(mockClientSession.closeGracefully()).thenReturn(Mono.empty());
|
||||
|
||||
initializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS,
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook);
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHook() {
|
||||
AtomicReference<Initialization> capturedInit = new AtomicReference<>();
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
capturedInit.set(invocation.getArgument(0));
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
// Verify hook was called
|
||||
verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
|
||||
|
||||
// Verify the hook received correct initialization data
|
||||
assertThat(capturedInit.get()).isNotNull();
|
||||
assertThat(capturedInit.get().mcpSession()).isEqualTo(mockClientSession);
|
||||
assertThat(capturedInit.get().initializeResult()).isEqualTo(MOCK_INIT_RESULT);
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHookOnlyOnce() {
|
||||
// First initialization
|
||||
StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
|
||||
.expectNext("result1")
|
||||
.verifyComplete();
|
||||
|
||||
// Second call should reuse initialization and NOT call hook again
|
||||
StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2")))
|
||||
.expectNext("result2")
|
||||
.verifyComplete();
|
||||
|
||||
// Hook should only be called once
|
||||
verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHookOnlyOnceWithConcurrentRequests() {
|
||||
AtomicInteger hookInvocationCount = new AtomicInteger(0);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
hookInvocationCount.incrementAndGet();
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
// Start multiple concurrent initializations
|
||||
Mono<String> init1 = initializer.withInitialization("test1", init -> Mono.just("result1"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
Mono<String> init2 = initializer.withInitialization("test2", init -> Mono.just("result2"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
Mono<String> init3 = initializer.withInitialization("test3", init -> Mono.just("result3"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
|
||||
// TODO: can we assume the order of results?
|
||||
StepVerifier.create(Mono.zip(init1, init2, init3)).assertNext(tuple -> {
|
||||
assertThat(tuple.getT1()).isEqualTo("result1");
|
||||
assertThat(tuple.getT2()).isEqualTo("result2");
|
||||
assertThat(tuple.getT3()).isEqualTo("result3");
|
||||
}).verifyComplete();
|
||||
|
||||
// Hook should only be called once despite concurrent requests
|
||||
assertThat(hookInvocationCount.get()).isEqualTo(1);
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldFailInitializationWhenPostInitializationHookFails() {
|
||||
RuntimeException hookError = new RuntimeException("Post-initialization hook failed");
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.error(hookError));
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectErrorMatches(ex -> ex instanceof RuntimeException && ex.getCause() == hookError)
|
||||
.verify();
|
||||
|
||||
// Verify initialization was not completed
|
||||
assertThat(initializer.isInitialized()).isFalse();
|
||||
assertThat(initializer.currentInitializationResult()).isNull();
|
||||
|
||||
// Verify the hook was called
|
||||
verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldNotInvokePostInitializationHookWhenInitializationFails() {
|
||||
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.error(new RuntimeException("Initialization failed")));
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
// Hook should NOT be called when initialization fails
|
||||
verify(mockPostInitializationHook, never()).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldNotInvokePostInitializationHookWhenNotificationFails() {
|
||||
when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
|
||||
.thenReturn(Mono.error(new RuntimeException("Notification failed")));
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
// Hook should NOT be called when notification fails
|
||||
verify(mockPostInitializationHook, never()).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHookAgainAfterReinitialization() {
|
||||
AtomicInteger hookInvocationCount = new AtomicInteger(0);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
hookInvocationCount.incrementAndGet();
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
// First initialization
|
||||
StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
|
||||
.expectNext("result1")
|
||||
.verifyComplete();
|
||||
|
||||
assertThat(hookInvocationCount.get()).isEqualTo(1);
|
||||
|
||||
// Simulate transport session exception to trigger re-initialization
|
||||
initializer.handleException(new McpTransportSessionNotFoundException("Session lost"));
|
||||
|
||||
// Hook should be called twice (once for each initialization)
|
||||
assertThat(hookInvocationCount.get()).isEqualTo(2);
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldAllowPostInitializationHookToPerformAsyncOperations() {
|
||||
AtomicInteger operationCount = new AtomicInteger(0);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class)))
|
||||
.thenReturn(Mono.fromRunnable(() -> operationCount.incrementAndGet()).then());
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
// Verify the async operation was executed
|
||||
assertThat(operationCount.get()).isEqualTo(1);
|
||||
verify(mockPostInitializationHook, times(1)).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldProvideCorrectInitializationDataToHook() {
|
||||
AtomicReference<McpClientSession> capturedSession = new AtomicReference<>();
|
||||
AtomicReference<McpSchema.InitializeResult> capturedResult = new AtomicReference<>();
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
Initialization init = invocation.getArgument(0);
|
||||
capturedSession.set(init.mcpSession());
|
||||
capturedResult.set(init.initializeResult());
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
// Verify the hook received the correct session and result
|
||||
assertThat(capturedSession.get()).isEqualTo(mockClientSession);
|
||||
assertThat(capturedResult.get()).isEqualTo(MOCK_INIT_RESULT);
|
||||
assertThat(capturedResult.get().protocolVersion()).isEqualTo("2.0.0");
|
||||
assertThat(capturedResult.get().serverInfo().name()).isEqualTo("test-server");
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInvokePostInitializationHookAfterSuccessfulInitialization() {
|
||||
AtomicReference<Boolean> notificationSent = new AtomicReference<>(false);
|
||||
AtomicReference<Boolean> hookCalledAfterNotification = new AtomicReference<>(false);
|
||||
|
||||
when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
|
||||
.thenAnswer(invocation -> {
|
||||
notificationSent.set(true);
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenAnswer(invocation -> {
|
||||
// Due to flatMap chaining in doInitialize, if the hook is called,
|
||||
// the notification must have been sent first
|
||||
hookCalledAfterNotification.set(notificationSent.get());
|
||||
return Mono.empty();
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
// Verify the hook was called and notification was already sent at that point
|
||||
assertThat(hookCalledAfterNotification.get()).isTrue();
|
||||
verify(mockClientSession).sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any());
|
||||
verify(mockPostInitializationHook).apply(any(Initialization.class));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
@ -10,14 +10,14 @@ import java.util.concurrent.atomic.AtomicInteger;
|
|||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
import io.modelcontextprotocol.client.LifecycleInitializer.Initialization;
|
||||
import io.modelcontextprotocol.spec.McpClientSession;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.MockitoAnnotations;
|
||||
|
||||
import io.modelcontextprotocol.spec.McpClientSession;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpTransportSessionNotFoundException;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
import reactor.test.StepVerifier;
|
||||
|
|
@ -58,12 +58,16 @@ class LifecycleInitializerTests {
|
|||
@Mock
|
||||
private Function<ContextView, McpClientSession> mockSessionSupplier;
|
||||
|
||||
@Mock
|
||||
private Function<Initialization, Mono<Void>> mockPostInitializationHook;
|
||||
|
||||
private LifecycleInitializer initializer;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
MockitoAnnotations.openMocks(this);
|
||||
|
||||
when(mockPostInitializationHook.apply(any(Initialization.class))).thenReturn(Mono.empty());
|
||||
when(mockSessionSupplier.apply(any(ContextView.class))).thenReturn(mockClientSession);
|
||||
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.just(MOCK_INIT_RESULT));
|
||||
|
|
@ -72,45 +76,45 @@ class LifecycleInitializerTests {
|
|||
when(mockClientSession.closeGracefully()).thenReturn(Mono.empty());
|
||||
|
||||
initializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS,
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier);
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook);
|
||||
}
|
||||
|
||||
@Test
|
||||
void constructorShouldValidateParameters() {
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(null, CLIENT_INFO, PROTOCOL_VERSIONS, INITIALIZATION_TIMEOUT,
|
||||
mockSessionSupplier))
|
||||
mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Client capabilities must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, null, PROTOCOL_VERSIONS,
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier))
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Client info must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, null,
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier))
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Protocol versions must not be empty");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, List.of(),
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier))
|
||||
INITIALIZATION_TIMEOUT, mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Protocol versions must not be empty");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS, null,
|
||||
mockSessionSupplier))
|
||||
mockSessionSupplier, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Initialization timeout must not be null");
|
||||
|
||||
assertThatThrownBy(() -> new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO, PROTOCOL_VERSIONS,
|
||||
INITIALIZATION_TIMEOUT, null))
|
||||
INITIALIZATION_TIMEOUT, null, mockPostInitializationHook))
|
||||
.isInstanceOf(IllegalArgumentException.class)
|
||||
.hasMessageContaining("Session supplier must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldInitializeSuccessfully() {
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.assertNext(result -> {
|
||||
assertThat(result).isEqualTo(MOCK_INIT_RESULT);
|
||||
assertThat(initializer.isInitialized()).isTrue();
|
||||
|
|
@ -132,7 +136,7 @@ class LifecycleInitializerTests {
|
|||
return Mono.just(MOCK_INIT_RESULT);
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.assertNext(result -> {
|
||||
assertThat(capturedRequest.get().protocolVersion()).isEqualTo("2.0.0"); // Latest
|
||||
// version
|
||||
|
|
@ -152,7 +156,7 @@ class LifecycleInitializerTests {
|
|||
when(mockClientSession.sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.just(unsupportedResult));
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
|
|
@ -167,13 +171,13 @@ class LifecycleInitializerTests {
|
|||
Duration SLOW_RESPONSE_DELAY = Duration.ofSeconds(5);
|
||||
|
||||
LifecycleInitializer shortTimeoutInitializer = new LifecycleInitializer(CLIENT_CAPABILITIES, CLIENT_INFO,
|
||||
PROTOCOL_VERSIONS, INITIALIZE_TIMEOUT, mockSessionSupplier);
|
||||
PROTOCOL_VERSIONS, INITIALIZE_TIMEOUT, mockSessionSupplier, mockPostInitializationHook);
|
||||
|
||||
when(mockClientSession.<McpSchema.InitializeResult>sendRequest(eq(McpSchema.METHOD_INITIALIZE), any(), any()))
|
||||
.thenReturn(Mono.just(MOCK_INIT_RESULT).delayElement(SLOW_RESPONSE_DELAY, virtualTimeScheduler));
|
||||
|
||||
StepVerifier
|
||||
.withVirtualTime(() -> shortTimeoutInitializer.withIntitialization("test",
|
||||
.withVirtualTime(() -> shortTimeoutInitializer.withInitialization("test",
|
||||
init -> Mono.just(init.initializeResult())), () -> virtualTimeScheduler, Long.MAX_VALUE)
|
||||
.expectSubscription()
|
||||
.expectNoEvent(INITIALIZE_TIMEOUT)
|
||||
|
|
@ -184,12 +188,12 @@ class LifecycleInitializerTests {
|
|||
@Test
|
||||
void shouldReuseExistingInitialization() {
|
||||
// First initialization
|
||||
StepVerifier.create(initializer.withIntitialization("test1", init -> Mono.just("result1")))
|
||||
StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
|
||||
.expectNext("result1")
|
||||
.verifyComplete();
|
||||
|
||||
// Second call should reuse the same initialization
|
||||
StepVerifier.create(initializer.withIntitialization("test2", init -> Mono.just("result2")))
|
||||
StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2")))
|
||||
.expectNext("result2")
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -209,11 +213,11 @@ class LifecycleInitializerTests {
|
|||
|
||||
// Start multiple concurrent initializations using subscribeOn with parallel
|
||||
// scheduler
|
||||
Mono<String> init1 = initializer.withIntitialization("test1", init -> Mono.just("result1"))
|
||||
Mono<String> init1 = initializer.withInitialization("test1", init -> Mono.just("result1"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
Mono<String> init2 = initializer.withIntitialization("test2", init -> Mono.just("result2"))
|
||||
Mono<String> init2 = initializer.withInitialization("test2", init -> Mono.just("result2"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
Mono<String> init3 = initializer.withIntitialization("test3", init -> Mono.just("result3"))
|
||||
Mono<String> init3 = initializer.withInitialization("test3", init -> Mono.just("result3"))
|
||||
.subscribeOn(Schedulers.parallel());
|
||||
|
||||
StepVerifier.create(Mono.zip(init1, init2, init3)).assertNext(tuple -> {
|
||||
|
|
@ -235,7 +239,7 @@ class LifecycleInitializerTests {
|
|||
// succeeds on the second call
|
||||
.thenReturn(Mono.just(MOCK_INIT_RESULT));
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
|
|
@ -244,7 +248,7 @@ class LifecycleInitializerTests {
|
|||
|
||||
// The initializer can recover from previous errors
|
||||
StepVerifier
|
||||
.create(initializer.withIntitialization("successful init", init -> Mono.just(init.initializeResult())))
|
||||
.create(initializer.withInitialization("successful init", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -255,7 +259,7 @@ class LifecycleInitializerTests {
|
|||
@Test
|
||||
void shouldHandleTransportSessionNotFoundException() {
|
||||
// successful initialization first
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -277,7 +281,7 @@ class LifecycleInitializerTests {
|
|||
@Test
|
||||
void shouldHandleOtherExceptions() {
|
||||
// Simulate a successful initialization first
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -295,7 +299,7 @@ class LifecycleInitializerTests {
|
|||
|
||||
@Test
|
||||
void shouldCloseGracefully() {
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -307,7 +311,7 @@ class LifecycleInitializerTests {
|
|||
|
||||
@Test
|
||||
void shouldCloseImmediately() {
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -342,7 +346,7 @@ class LifecycleInitializerTests {
|
|||
new McpSchema.Implementation("test-server", "1.0.0"), "Test instructions"));
|
||||
});
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.assertNext(result -> {
|
||||
// Latest from new versions
|
||||
assertThat(capturedRequest.get().protocolVersion()).isEqualTo("4.0.0");
|
||||
|
|
@ -363,7 +367,7 @@ class LifecycleInitializerTests {
|
|||
});
|
||||
|
||||
StepVerifier
|
||||
.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult()))
|
||||
.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult()))
|
||||
.contextWrite(Context.of(contextKey, contextValue)))
|
||||
.expectNext(MOCK_INIT_RESULT)
|
||||
.verifyComplete();
|
||||
|
|
@ -374,7 +378,7 @@ class LifecycleInitializerTests {
|
|||
|
||||
@Test
|
||||
void shouldProvideAccessToMcpSessionAndInitializeResult() {
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> {
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> {
|
||||
assertThat(init.mcpSession()).isEqualTo(mockClientSession);
|
||||
assertThat(init.initializeResult()).isEqualTo(MOCK_INIT_RESULT);
|
||||
return Mono.just("success");
|
||||
|
|
@ -386,7 +390,7 @@ class LifecycleInitializerTests {
|
|||
when(mockClientSession.sendNotification(eq(McpSchema.METHOD_NOTIFICATION_INITIALIZED), any()))
|
||||
.thenReturn(Mono.error(new RuntimeException("Notification failed")));
|
||||
|
||||
StepVerifier.create(initializer.withIntitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
StepVerifier.create(initializer.withInitialization("test", init -> Mono.just(init.initializeResult())))
|
||||
.expectError(RuntimeException.class)
|
||||
.verify();
|
||||
|
||||
|
|
@ -403,7 +407,7 @@ class LifecycleInitializerTests {
|
|||
@Test
|
||||
void shouldReinitializeAfterTransportSessionException() {
|
||||
// First initialization
|
||||
StepVerifier.create(initializer.withIntitialization("test1", init -> Mono.just("result1")))
|
||||
StepVerifier.create(initializer.withInitialization("test1", init -> Mono.just("result1")))
|
||||
.expectNext("result1")
|
||||
.verifyComplete();
|
||||
|
||||
|
|
@ -411,7 +415,7 @@ class LifecycleInitializerTests {
|
|||
initializer.handleException(new McpTransportSessionNotFoundException("Session lost"));
|
||||
|
||||
// Should be able to initialize again
|
||||
StepVerifier.create(initializer.withIntitialization("test2", init -> Mono.just("result2")))
|
||||
StepVerifier.create(initializer.withInitialization("test2", init -> Mono.just("result2")))
|
||||
.expectNext("result2")
|
||||
.verifyComplete();
|
||||
|
||||
|
|
|
|||
|
|
@ -10,12 +10,20 @@ import io.modelcontextprotocol.spec.McpSchema;
|
|||
import io.modelcontextprotocol.spec.ProtocolVersions;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.core.publisher.Mono;
|
||||
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.Objects;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
import java.util.function.Function;
|
||||
|
||||
import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatCode;
|
||||
|
||||
class McpAsyncClientTests {
|
||||
|
|
@ -24,6 +32,7 @@ class McpAsyncClientTests {
|
|||
"1.0.0");
|
||||
|
||||
public static final McpSchema.ServerCapabilities MOCK_SERVER_CAPABILITIES = McpSchema.ServerCapabilities.builder()
|
||||
.tools(true)
|
||||
.build();
|
||||
|
||||
public static final McpSchema.InitializeResult MOCK_INIT_RESULT = new McpSchema.InitializeResult(
|
||||
|
|
@ -31,6 +40,91 @@ class McpAsyncClientTests {
|
|||
|
||||
private static final String CONTEXT_KEY = "context.key";
|
||||
|
||||
private McpClientTransport createMockTransportForToolValidation(boolean hasOutputSchema, boolean invalidOutput)
|
||||
throws JsonProcessingException {
|
||||
|
||||
// Create tool with or without output schema
|
||||
Map<String, Object> inputSchemaMap = Map.of("type", "object", "properties",
|
||||
Map.of("expression", Map.of("type", "string")), "required", List.of("expression"));
|
||||
|
||||
McpSchema.JsonSchema inputSchema = new McpSchema.JsonSchema("object", inputSchemaMap, null, null, null, null);
|
||||
McpSchema.Tool.Builder toolBuilder = McpSchema.Tool.builder()
|
||||
.name("calculator")
|
||||
.description("Performs mathematical calculations")
|
||||
.inputSchema(inputSchema);
|
||||
|
||||
if (hasOutputSchema) {
|
||||
Map<String, Object> outputSchema = Map.of("type", "object", "properties",
|
||||
Map.of("result", Map.of("type", "number"), "operation", Map.of("type", "string")), "required",
|
||||
List.of("result", "operation"));
|
||||
toolBuilder.outputSchema(outputSchema);
|
||||
}
|
||||
|
||||
McpSchema.Tool calculatorTool = toolBuilder.build();
|
||||
McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(calculatorTool), null);
|
||||
|
||||
// Create call tool result - valid or invalid based on parameter
|
||||
Map<String, Object> structuredContent = invalidOutput ? Map.of("result", "5", "operation", "add")
|
||||
: Map.of("result", 5, "operation", "add");
|
||||
|
||||
McpSchema.CallToolResult mockCallToolResult = McpSchema.CallToolResult.builder()
|
||||
.addTextContent("Calculation result")
|
||||
.structuredContent(structuredContent)
|
||||
.build();
|
||||
|
||||
return new McpClientTransport() {
|
||||
Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler;
|
||||
|
||||
@Override
|
||||
public Mono<Void> connect(
|
||||
Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
|
||||
this.handler = handler;
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> closeGracefully() {
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
|
||||
if (!(message instanceof McpSchema.JSONRPCRequest request)) {
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
McpSchema.JSONRPCResponse response;
|
||||
if (McpSchema.METHOD_INITIALIZE.equals(request.method())) {
|
||||
response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), MOCK_INIT_RESULT,
|
||||
null);
|
||||
}
|
||||
else if (McpSchema.METHOD_TOOLS_LIST.equals(request.method())) {
|
||||
response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), mockToolsResult,
|
||||
null);
|
||||
}
|
||||
else if (McpSchema.METHOD_TOOLS_CALL.equals(request.method())) {
|
||||
response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
mockCallToolResult, null);
|
||||
}
|
||||
else {
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
return handler.apply(Mono.just(response)).then();
|
||||
}
|
||||
|
||||
@Override
|
||||
public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
|
||||
return JSON_MAPPER.convertValue(data, new TypeRef<>() {
|
||||
@Override
|
||||
public java.lang.reflect.Type getType() {
|
||||
return typeRef.getType();
|
||||
}
|
||||
});
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
@Test
|
||||
void validateContextPassedToTransportConnect() {
|
||||
McpClientTransport transport = new McpClientTransport() {
|
||||
|
|
@ -87,4 +181,64 @@ class McpAsyncClientTests {
|
|||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallToolWithOutputSchemaValidationSuccess() throws JsonProcessingException {
|
||||
McpClientTransport transport = createMockTransportForToolValidation(true, false);
|
||||
|
||||
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
|
||||
|
||||
StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
|
||||
|
||||
StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))))
|
||||
.expectNextMatches(response -> {
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response.isError()).isFalse();
|
||||
assertThat(response.structuredContent()).isInstanceOf(Map.class);
|
||||
assertThat((Map<?, ?>) response.structuredContent()).hasSize(2);
|
||||
assertThat(response.content()).hasSize(1);
|
||||
return true;
|
||||
})
|
||||
.verifyComplete();
|
||||
|
||||
StepVerifier.create(client.closeGracefully()).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallToolWithNoOutputSchemaSuccess() throws JsonProcessingException {
|
||||
McpClientTransport transport = createMockTransportForToolValidation(false, false);
|
||||
|
||||
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
|
||||
|
||||
StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
|
||||
|
||||
StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))))
|
||||
.expectNextMatches(response -> {
|
||||
assertThat(response).isNotNull();
|
||||
assertThat(response.isError()).isFalse();
|
||||
assertThat(response.structuredContent()).isInstanceOf(Map.class);
|
||||
assertThat((Map<?, ?>) response.structuredContent()).hasSize(2);
|
||||
assertThat(response.content()).hasSize(1);
|
||||
return true;
|
||||
})
|
||||
.verifyComplete();
|
||||
|
||||
StepVerifier.create(client.closeGracefully()).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallToolWithOutputSchemaValidationFailure() throws JsonProcessingException {
|
||||
McpClientTransport transport = createMockTransportForToolValidation(true, true);
|
||||
|
||||
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
|
||||
|
||||
StepVerifier.create(client.initialize()).expectNextMatches(Objects::nonNull).verifyComplete();
|
||||
|
||||
StepVerifier.create(client.callTool(new McpSchema.CallToolRequest("calculator", Map.of("expression", "2 + 3"))))
|
||||
.expectErrorMatches(ex -> ex instanceof IllegalArgumentException
|
||||
&& ex.getMessage().contains("Tool call result validation failed"))
|
||||
.verify();
|
||||
|
||||
StepVerifier.create(client.closeGracefully()).verifyComplete();
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,9 +8,9 @@ import java.time.Duration;
|
|||
import java.util.List;
|
||||
|
||||
import io.modelcontextprotocol.MockMcpClientTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.test.StepVerifier;
|
||||
|
|
@ -22,7 +22,7 @@ import static org.assertj.core.api.Assertions.assertThat;
|
|||
*/
|
||||
class McpClientProtocolVersionTests {
|
||||
|
||||
private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(30);
|
||||
private static final Duration REQUEST_TIMEOUT = Duration.ofSeconds(300);
|
||||
|
||||
private static final McpSchema.Implementation CLIENT_INFO = new McpSchema.Implementation("test-client", "1.0.0");
|
||||
|
||||
|
|
@ -46,13 +46,12 @@ class McpClientProtocolVersionTests {
|
|||
assertThat(initRequest.protocolVersion()).isEqualTo(transport.protocolVersions().get(0));
|
||||
|
||||
transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
new McpSchema.InitializeResult(protocolVersion, null,
|
||||
new McpSchema.InitializeResult(protocolVersion, ServerCapabilities.builder().build(),
|
||||
new McpSchema.Implementation("test-server", "1.0.0"), null),
|
||||
null));
|
||||
}).assertNext(result -> {
|
||||
assertThat(result.protocolVersion()).isEqualTo(protocolVersion);
|
||||
}).verifyComplete();
|
||||
|
||||
}
|
||||
finally {
|
||||
// Ensure cleanup happens even if test fails
|
||||
|
|
@ -81,7 +80,7 @@ class McpClientProtocolVersionTests {
|
|||
assertThat(initRequest.protocolVersion()).isIn(List.of(oldVersion, McpSchema.LATEST_PROTOCOL_VERSION));
|
||||
|
||||
transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
new McpSchema.InitializeResult(oldVersion, null,
|
||||
new McpSchema.InitializeResult(oldVersion, ServerCapabilities.builder().build(),
|
||||
new McpSchema.Implementation("test-server", "1.0.0"), null),
|
||||
null));
|
||||
}).assertNext(result -> {
|
||||
|
|
@ -110,7 +109,7 @@ class McpClientProtocolVersionTests {
|
|||
assertThat(request.params()).isInstanceOf(McpSchema.InitializeRequest.class);
|
||||
|
||||
transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
new McpSchema.InitializeResult(unsupportedVersion, null,
|
||||
new McpSchema.InitializeResult(unsupportedVersion, ServerCapabilities.builder().build(),
|
||||
new McpSchema.Implementation("test-server", "1.0.0"), null),
|
||||
null));
|
||||
}).expectError(RuntimeException.class).verify();
|
||||
|
|
@ -143,7 +142,7 @@ class McpClientProtocolVersionTests {
|
|||
assertThat(initRequest.protocolVersion()).isEqualTo(latestVersion);
|
||||
|
||||
transport.simulateIncomingMessage(new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(),
|
||||
new McpSchema.InitializeResult(latestVersion, null,
|
||||
new McpSchema.InitializeResult(latestVersion, ServerCapabilities.builder().build(),
|
||||
new McpSchema.Implementation("test-server", "1.0.0"), null),
|
||||
null));
|
||||
}).assertNext(result -> {
|
||||
|
|
|
|||
Loading…
Reference in a new issue