refactor(client): improve validation and remove server methods (#14)
Add client initialization and capability validation checks - New isInitialized() method to check client state - Validate server capabilities before tool/resource operations - Add clear error messages for common failure cases - Remove server-side notification methods from client: sendResourcesListChanged(), promptListChangedNotification() - Improve protocol version handling - Testing improvements and new initialization tests - Redesign MockMcpTransport internals and adapt tests - Correct McpError messages - clean unused imports Resolves #13 Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com> Signed-off-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com> Co-authored-by: Dariusz Jędrzejczyk <dariusz.jedrzejczyk@broadcom.com>
This commit is contained in:
Christian Tzolov
GitHub
parent
d9c4818607
commit
718042644a
4 changed files with 184 additions and 112 deletions
|
|
@ -4,7 +4,9 @@
|
|||
|
||||
package io.modelcontextprotocol;
|
||||
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.function.BiConsumer;
|
||||
import java.util.function.Function;
|
||||
|
||||
import com.fasterxml.jackson.core.type.TypeReference;
|
||||
|
|
@ -14,10 +16,8 @@ import io.modelcontextprotocol.spec.McpSchema;
|
|||
import io.modelcontextprotocol.spec.ServerMcpTransport;
|
||||
import io.modelcontextprotocol.spec.McpSchema.JSONRPCNotification;
|
||||
import io.modelcontextprotocol.spec.McpSchema.JSONRPCRequest;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.publisher.Sinks;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
/**
|
||||
* A mock implementation of the {@link ClientMcpTransport} and {@link ServerMcpTransport}
|
||||
|
|
@ -25,39 +25,44 @@ import reactor.core.scheduler.Schedulers;
|
|||
*/
|
||||
public class MockMcpTransport implements ClientMcpTransport, ServerMcpTransport {
|
||||
|
||||
private final AtomicInteger inboundMessageCount = new AtomicInteger(0);
|
||||
|
||||
private final Sinks.Many<McpSchema.JSONRPCMessage> outgoing = Sinks.many().multicast().onBackpressureBuffer();
|
||||
|
||||
private final Sinks.Many<McpSchema.JSONRPCMessage> inbound = Sinks.many().unicast().onBackpressureBuffer();
|
||||
|
||||
private final Flux<McpSchema.JSONRPCMessage> outboundView = outgoing.asFlux().cache(1);
|
||||
private final List<McpSchema.JSONRPCMessage> sent = new ArrayList<>();
|
||||
|
||||
private final BiConsumer<MockMcpTransport, McpSchema.JSONRPCMessage> interceptor;
|
||||
|
||||
public MockMcpTransport() {
|
||||
this((t, msg) -> {
|
||||
});
|
||||
}
|
||||
|
||||
public MockMcpTransport(BiConsumer<MockMcpTransport, McpSchema.JSONRPCMessage> interceptor) {
|
||||
this.interceptor = interceptor;
|
||||
}
|
||||
|
||||
public void simulateIncomingMessage(McpSchema.JSONRPCMessage message) {
|
||||
if (inbound.tryEmitNext(message).isFailure()) {
|
||||
throw new RuntimeException("Failed to emit message " + message);
|
||||
throw new RuntimeException("Failed to process incoming message " + message);
|
||||
}
|
||||
inboundMessageCount.incrementAndGet();
|
||||
}
|
||||
|
||||
@Override
|
||||
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
|
||||
if (outgoing.tryEmitNext(message).isFailure()) {
|
||||
return Mono.error(new RuntimeException("Can't emit outgoing message " + message));
|
||||
}
|
||||
sent.add(message);
|
||||
interceptor.accept(this, message);
|
||||
return Mono.empty();
|
||||
}
|
||||
|
||||
public McpSchema.JSONRPCRequest getLastSentMessageAsRequest() {
|
||||
return (JSONRPCRequest) outboundView.blockFirst();
|
||||
return (JSONRPCRequest) getLastSentMessage();
|
||||
}
|
||||
|
||||
public McpSchema.JSONRPCNotification getLastSentMessageAsNotifiation() {
|
||||
return (JSONRPCNotification) outboundView.blockFirst();
|
||||
public McpSchema.JSONRPCNotification getLastSentMessageAsNotification() {
|
||||
return (JSONRPCNotification) getLastSentMessage();
|
||||
}
|
||||
|
||||
public McpSchema.JSONRPCMessage getLastSentMessage() {
|
||||
return outboundView.blockFirst();
|
||||
return !sent.isEmpty() ? sent.get(sent.size() - 1) : null;
|
||||
}
|
||||
|
||||
private volatile boolean connected = false;
|
||||
|
|
@ -69,7 +74,6 @@ public class MockMcpTransport implements ClientMcpTransport, ServerMcpTransport
|
|||
}
|
||||
connected = true;
|
||||
return inbound.asFlux()
|
||||
.publishOn(Schedulers.boundedElastic())
|
||||
.flatMap(message -> Mono.just(message).transform(handler))
|
||||
.doFinally(signal -> connected = false)
|
||||
.then();
|
||||
|
|
@ -79,7 +83,6 @@ public class MockMcpTransport implements ClientMcpTransport, ServerMcpTransport
|
|||
public Mono<Void> closeGracefully() {
|
||||
return Mono.defer(() -> {
|
||||
connected = false;
|
||||
outgoing.tryEmitComplete();
|
||||
inbound.tryEmitComplete();
|
||||
// Wait for all subscribers to complete
|
||||
return Mono.empty();
|
||||
|
|
|
|||
|
|
@ -275,9 +275,6 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
assertThatCode(() -> {
|
||||
client.initialize().block();
|
||||
// Trigger notifications
|
||||
client.sendResourcesListChanged().block();
|
||||
client.promptListChangedNotification().block();
|
||||
client.closeGracefully().block();
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -259,9 +259,6 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
assertThatCode(() -> {
|
||||
client.initialize();
|
||||
// Trigger notifications
|
||||
client.sendResourcesListChanged();
|
||||
client.promptListChangedNotification();
|
||||
client.close();
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@
|
|||
|
||||
package io.modelcontextprotocol.client;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
|
@ -17,19 +16,82 @@ import io.modelcontextprotocol.MockMcpTransport;
|
|||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
|
||||
import io.modelcontextprotocol.spec.McpSchema.InitializeResult;
|
||||
import io.modelcontextprotocol.spec.McpSchema.Root;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import reactor.core.publisher.Mono;
|
||||
|
||||
import static io.modelcontextprotocol.spec.McpSchema.METHOD_INITIALIZE;
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
import static org.assertj.core.api.Assertions.assertThatThrownBy;
|
||||
import static org.awaitility.Awaitility.await;
|
||||
|
||||
class McpAsyncClientResponseHandlerTests {
|
||||
|
||||
private static final McpSchema.Implementation SERVER_INFO = new McpSchema.Implementation("test-server", "1.0.0");
|
||||
|
||||
private static final McpSchema.ServerCapabilities SERVER_CAPABILITIES = McpSchema.ServerCapabilities.builder()
|
||||
.tools(true)
|
||||
.resources(true, true) // Enable both resources and resource templates
|
||||
.build();
|
||||
|
||||
private static MockMcpTransport initializationEnabledTransport() {
|
||||
return initializationEnabledTransport(SERVER_CAPABILITIES, SERVER_INFO);
|
||||
}
|
||||
|
||||
private static MockMcpTransport initializationEnabledTransport(McpSchema.ServerCapabilities mockServerCapabilities,
|
||||
McpSchema.Implementation mockServerInfo) {
|
||||
McpSchema.InitializeResult mockInitResult = new McpSchema.InitializeResult(McpSchema.LATEST_PROTOCOL_VERSION,
|
||||
mockServerCapabilities, mockServerInfo, "Test instructions");
|
||||
|
||||
return new MockMcpTransport((t, message) -> {
|
||||
if (message instanceof McpSchema.JSONRPCRequest r && METHOD_INITIALIZE.equals(r.method())) {
|
||||
McpSchema.JSONRPCResponse initResponse = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION,
|
||||
r.id(), mockInitResult, null);
|
||||
t.simulateIncomingMessage(initResponse);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSuccessfulInitialization() {
|
||||
McpSchema.Implementation serverInfo = new McpSchema.Implementation("mcp-test-server", "0.0.1");
|
||||
McpSchema.ServerCapabilities serverCapabilities = McpSchema.ServerCapabilities.builder()
|
||||
.tools(false)
|
||||
.resources(true, true) // Enable both resources and resource templates
|
||||
.build();
|
||||
MockMcpTransport transport = initializationEnabledTransport(serverCapabilities, serverInfo);
|
||||
McpAsyncClient asyncMcpClient = McpClient.async(transport).build();
|
||||
|
||||
// Verify client is not initialized initially
|
||||
assertThat(asyncMcpClient.isInitialized()).isFalse();
|
||||
|
||||
// Start initialization with reactive handling
|
||||
InitializeResult result = asyncMcpClient.initialize().block();
|
||||
|
||||
// Verify initialized notification was sent
|
||||
McpSchema.JSONRPCMessage notificationMessage = transport.getLastSentMessage();
|
||||
assertThat(notificationMessage).isInstanceOf(McpSchema.JSONRPCNotification.class);
|
||||
McpSchema.JSONRPCNotification notification = (McpSchema.JSONRPCNotification) notificationMessage;
|
||||
assertThat(notification.method()).isEqualTo(McpSchema.METHOD_NOTIFICATION_INITIALIZED);
|
||||
|
||||
// Verify initialization result
|
||||
assertThat(result).isNotNull();
|
||||
assertThat(result.protocolVersion()).isEqualTo(McpSchema.LATEST_PROTOCOL_VERSION);
|
||||
assertThat(result.capabilities()).isEqualTo(serverCapabilities);
|
||||
assertThat(result.serverInfo()).isEqualTo(serverInfo);
|
||||
assertThat(result.instructions()).isEqualTo("Test instructions");
|
||||
|
||||
// Verify client state after initialization
|
||||
assertThat(asyncMcpClient.isInitialized()).isTrue();
|
||||
assertThat(asyncMcpClient.getServerCapabilities()).isEqualTo(serverCapabilities);
|
||||
assertThat(asyncMcpClient.getServerInfo()).isEqualTo(serverInfo);
|
||||
|
||||
asyncMcpClient.closeGracefully();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testToolsChangeNotificationHandling() throws JsonProcessingException {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
MockMcpTransport transport = initializationEnabledTransport();
|
||||
|
||||
// Create a list to store received tools for verification
|
||||
List<McpSchema.Tool> receivedTools = new ArrayList<>();
|
||||
|
|
@ -41,6 +103,8 @@ class McpAsyncClientResponseHandlerTests {
|
|||
// Create client with tools change consumer
|
||||
McpAsyncClient asyncMcpClient = McpClient.async(transport).toolsChangeConsumer(toolsChangeConsumer).build();
|
||||
|
||||
assertThat(asyncMcpClient.initialize().block()).isNotNull();
|
||||
|
||||
// Create a mock tools list that the server will return
|
||||
Map<String, Object> inputSchema = Map.of("type", "object", "properties", Map.of(), "required", List.of());
|
||||
McpSchema.Tool mockTool = new McpSchema.Tool("test-tool", "Test Tool Description",
|
||||
|
|
@ -61,23 +125,23 @@ class McpAsyncClientResponseHandlerTests {
|
|||
transport.simulateIncomingMessage(toolsListResponse);
|
||||
|
||||
// Verify the consumer received the expected tools
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(receivedTools).hasSize(1);
|
||||
assertThat(receivedTools.get(0).name()).isEqualTo("test-tool");
|
||||
assertThat(receivedTools.get(0).description()).isEqualTo("Test Tool Description");
|
||||
});
|
||||
assertThat(receivedTools).hasSize(1);
|
||||
assertThat(receivedTools.get(0).name()).isEqualTo("test-tool");
|
||||
assertThat(receivedTools.get(0).description()).isEqualTo("Test Tool Description");
|
||||
|
||||
asyncMcpClient.closeGracefully();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRootsListRequestHandling() {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
MockMcpTransport transport = initializationEnabledTransport();
|
||||
|
||||
McpAsyncClient asyncMcpClient = McpClient.async(transport)
|
||||
.roots(new Root("file:///test/path", "test-root"))
|
||||
.build();
|
||||
|
||||
assertThat(asyncMcpClient.initialize().block()).isNotNull();
|
||||
|
||||
// Simulate incoming request
|
||||
McpSchema.JSONRPCRequest request = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION,
|
||||
McpSchema.METHOD_ROOTS_LIST, "test-id", null);
|
||||
|
|
@ -98,7 +162,7 @@ class McpAsyncClientResponseHandlerTests {
|
|||
|
||||
@Test
|
||||
void testResourcesChangeNotificationHandling() {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
MockMcpTransport transport = initializationEnabledTransport();
|
||||
|
||||
// Create a list to store received resources for verification
|
||||
List<McpSchema.Resource> receivedResources = new ArrayList<>();
|
||||
|
|
@ -112,6 +176,8 @@ class McpAsyncClientResponseHandlerTests {
|
|||
.resourcesChangeConsumer(resourcesChangeConsumer)
|
||||
.build();
|
||||
|
||||
assertThat(asyncMcpClient.initialize().block()).isNotNull();
|
||||
|
||||
// Create a mock resources list that the server will return
|
||||
McpSchema.Resource mockResource = new McpSchema.Resource("test://resource", "Test Resource", "A test resource",
|
||||
"text/plain", null);
|
||||
|
|
@ -132,19 +198,17 @@ class McpAsyncClientResponseHandlerTests {
|
|||
transport.simulateIncomingMessage(resourcesListResponse);
|
||||
|
||||
// Verify the consumer received the expected resources
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(receivedResources).hasSize(1);
|
||||
assertThat(receivedResources.get(0).uri()).isEqualTo("test://resource");
|
||||
assertThat(receivedResources.get(0).name()).isEqualTo("Test Resource");
|
||||
assertThat(receivedResources.get(0).description()).isEqualTo("A test resource");
|
||||
});
|
||||
assertThat(receivedResources).hasSize(1);
|
||||
assertThat(receivedResources.get(0).uri()).isEqualTo("test://resource");
|
||||
assertThat(receivedResources.get(0).name()).isEqualTo("Test Resource");
|
||||
assertThat(receivedResources.get(0).description()).isEqualTo("A test resource");
|
||||
|
||||
asyncMcpClient.closeGracefully();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testPromptsChangeNotificationHandling() {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
MockMcpTransport transport = initializationEnabledTransport();
|
||||
|
||||
// Create a list to store received prompts for verification
|
||||
List<McpSchema.Prompt> receivedPrompts = new ArrayList<>();
|
||||
|
|
@ -156,6 +220,8 @@ class McpAsyncClientResponseHandlerTests {
|
|||
// Create client with prompts change consumer
|
||||
McpAsyncClient asyncMcpClient = McpClient.async(transport).promptsChangeConsumer(promptsChangeConsumer).build();
|
||||
|
||||
assertThat(asyncMcpClient.initialize().block()).isNotNull();
|
||||
|
||||
// Create a mock prompts list that the server will return
|
||||
McpSchema.Prompt mockPrompt = new McpSchema.Prompt("test-prompt", "Test Prompt Description",
|
||||
List.of(new McpSchema.PromptArgument("arg1", "Test argument", true)));
|
||||
|
|
@ -175,20 +241,18 @@ class McpAsyncClientResponseHandlerTests {
|
|||
transport.simulateIncomingMessage(promptsListResponse);
|
||||
|
||||
// Verify the consumer received the expected prompts
|
||||
await().atMost(Duration.ofSeconds(5)).untilAsserted(() -> {
|
||||
assertThat(receivedPrompts).hasSize(1);
|
||||
assertThat(receivedPrompts.get(0).name()).isEqualTo("test-prompt");
|
||||
assertThat(receivedPrompts.get(0).description()).isEqualTo("Test Prompt Description");
|
||||
assertThat(receivedPrompts.get(0).arguments()).hasSize(1);
|
||||
assertThat(receivedPrompts.get(0).arguments().get(0).name()).isEqualTo("arg1");
|
||||
});
|
||||
assertThat(receivedPrompts).hasSize(1);
|
||||
assertThat(receivedPrompts.get(0).name()).isEqualTo("test-prompt");
|
||||
assertThat(receivedPrompts.get(0).description()).isEqualTo("Test Prompt Description");
|
||||
assertThat(receivedPrompts.get(0).arguments()).hasSize(1);
|
||||
assertThat(receivedPrompts.get(0).arguments().get(0).name()).isEqualTo("arg1");
|
||||
|
||||
asyncMcpClient.closeGracefully();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSamplingCreateMessageRequestHandling() {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
MockMcpTransport transport = initializationEnabledTransport();
|
||||
|
||||
// Create a test sampling handler that echoes back the input
|
||||
Function<McpSchema.CreateMessageRequest, Mono<McpSchema.CreateMessageResult>> samplingHandler = request -> {
|
||||
|
|
@ -203,6 +267,8 @@ class McpAsyncClientResponseHandlerTests {
|
|||
.sampling(samplingHandler)
|
||||
.build();
|
||||
|
||||
assertThat(asyncMcpClient.initialize().block()).isNotNull();
|
||||
|
||||
// Create a mock create message request
|
||||
var messageRequest = new McpSchema.CreateMessageRequest(
|
||||
List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message"))),
|
||||
|
|
@ -240,13 +306,15 @@ class McpAsyncClientResponseHandlerTests {
|
|||
|
||||
@Test
|
||||
void testSamplingCreateMessageRequestHandlingWithoutCapability() {
|
||||
MockMcpTransport transport = new MockMcpTransport();
|
||||
MockMcpTransport transport = initializationEnabledTransport();
|
||||
|
||||
// Create client without sampling capability
|
||||
McpAsyncClient asyncMcpClient = McpClient.async(transport)
|
||||
.capabilities(ClientCapabilities.builder().build()) // No sampling capability
|
||||
.build();
|
||||
|
||||
assertThat(asyncMcpClient.initialize().block()).isNotNull();
|
||||
|
||||
// Create a mock create message request
|
||||
var messageRequest = new McpSchema.CreateMessageRequest(
|
||||
List.of(new McpSchema.SamplingMessage(McpSchema.Role.USER, new McpSchema.TextContent("Test message"))),
|
||||
|
|
|
|||
Loading…
Reference in a new issue