feat(client): Improve initialization state handling in McpAsyncClient (#39)
- Add proper initialization state tracking using AtomicBoolean and Sinks - Implement timeout handling for requests requiring initialization - Ensure all client methods verify initialization state before proceeding - Fix rootsListChangedNotification to check initialization state - Improve error messages for uninitialized client operations - improve JavaDoc - Add tests to verify proper error handling for uninitialized clients - Replace hardcoded timeout constants with configurable getTimeoutDuration() method - Remove automatic initialization in setUp methods to allow explicit testing Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
Christian Tzolov
GitHub
parent
4eff00c8da
commit
109aab20e8
2 changed files with 613 additions and 242 deletions
|
|
@ -10,6 +10,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
|
|||
import java.util.function.Function;
|
||||
|
||||
import io.modelcontextprotocol.spec.ClientMcpTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
|
||||
import io.modelcontextprotocol.spec.McpSchema.ClientCapabilities;
|
||||
|
|
@ -47,8 +48,6 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
protected ClientMcpTransport mcpTransport;
|
||||
|
||||
private static final Duration TIMEOUT = Duration.ofSeconds(20);
|
||||
|
||||
private static final String ECHO_TEST_MESSAGE = "Hello MCP Spring AI!";
|
||||
|
||||
abstract protected ClientMcpTransport createMcpTransport();
|
||||
|
|
@ -59,6 +58,10 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
protected void onClose() {
|
||||
}
|
||||
|
||||
protected Duration getTimeoutDuration() {
|
||||
return Duration.ofSeconds(2);
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
onStart();
|
||||
|
|
@ -66,10 +69,9 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
assertThatCode(() -> {
|
||||
mcpAsyncClient = McpClient.async(mcpTransport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.capabilities(ClientCapabilities.builder().roots(true).build())
|
||||
.build();
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
|
|
@ -92,8 +94,16 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
.hasMessage("Request timeout must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListToolsWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpAsyncClient.listTools(null).block()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before listing tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListTools() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
|
||||
StepVerifier.create(mcpAsyncClient.listTools(null)).consumeNextWith(result -> {
|
||||
assertThat(result.tools()).isNotNull().isNotEmpty();
|
||||
|
||||
|
|
@ -103,13 +113,30 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
}).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testPingWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpAsyncClient.ping().block()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before pinging the server");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testPing() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
assertThatCode(() -> mcpAsyncClient.ping().block()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallToolWithoutInitialization() {
|
||||
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", ECHO_TEST_MESSAGE));
|
||||
|
||||
assertThatThrownBy(() -> mcpAsyncClient.callTool(callToolRequest).block()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before calling tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallTool() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
|
||||
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", ECHO_TEST_MESSAGE));
|
||||
|
||||
StepVerifier.create(mcpAsyncClient.callTool(callToolRequest)).consumeNextWith(callToolResult -> {
|
||||
|
|
@ -122,13 +149,23 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
@Test
|
||||
void testCallToolWithInvalidTool() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
|
||||
CallToolRequest invalidRequest = new CallToolRequest("nonexistent_tool", Map.of("message", ECHO_TEST_MESSAGE));
|
||||
|
||||
assertThatThrownBy(() -> mcpAsyncClient.callTool(invalidRequest).block()).isInstanceOf(Exception.class);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResourcesWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpAsyncClient.listResources(null).block()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before listing resources");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResources() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
|
||||
StepVerifier.create(mcpAsyncClient.listResources(null)).consumeNextWith(resources -> {
|
||||
assertThat(resources).isNotNull().satisfies(result -> {
|
||||
assertThat(result.resources()).isNotNull();
|
||||
|
|
@ -147,8 +184,16 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
assertThat(mcpAsyncClient).isNotNull();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListPromptsWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpAsyncClient.listPrompts(null).block()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before listing prompts");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListPrompts() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
|
||||
StepVerifier.create(mcpAsyncClient.listPrompts(null)).consumeNextWith(prompts -> {
|
||||
assertThat(prompts).isNotNull().satisfies(result -> {
|
||||
assertThat(result.prompts()).isNotNull();
|
||||
|
|
@ -162,8 +207,18 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
}).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGetPromptWithoutInitialization() {
|
||||
GetPromptRequest request = new GetPromptRequest("simple_prompt", Map.of());
|
||||
|
||||
assertThatThrownBy(() -> mcpAsyncClient.getPrompt(request).block()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before getting prompts");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGetPrompt() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
|
||||
StepVerifier.create(mcpAsyncClient.getPrompt(new GetPromptRequest("simple_prompt", Map.of())))
|
||||
.consumeNextWith(prompt -> {
|
||||
assertThat(prompt).isNotNull().satisfies(result -> {
|
||||
|
|
@ -174,8 +229,16 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
.verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRootsListChangedWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpAsyncClient.rootsListChangedNotification().block()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before sending roots list changed notification");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRootsListChanged() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
|
||||
assertThatCode(() -> mcpAsyncClient.rootsListChangedNotification().block()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
|
|
@ -184,7 +247,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
var transport = createMcpTransport();
|
||||
|
||||
var client = McpClient.async(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.roots(new Root("file:///test/path", "test-root"))
|
||||
.build();
|
||||
|
||||
|
|
@ -233,8 +296,16 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
}).verifyComplete();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResourceTemplatesWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpAsyncClient.listResourceTemplates().block()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before listing resource templates");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResourceTemplates() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
|
||||
StepVerifier.create(mcpAsyncClient.listResourceTemplates()).consumeNextWith(result -> {
|
||||
assertThat(result).isNotNull();
|
||||
assertThat(result.resourceTemplates()).isNotNull();
|
||||
|
|
@ -266,7 +337,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
|
||||
var transport = createMcpTransport();
|
||||
var client = McpClient.async(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.toolsChangeConsumer(tools -> Mono.fromRunnable(() -> toolsNotificationReceived.set(true)))
|
||||
.resourcesChangeConsumer(resources -> Mono.fromRunnable(() -> resourcesNotificationReceived.set(true)))
|
||||
.promptsChangeConsumer(prompts -> Mono.fromRunnable(() -> promptsNotificationReceived.set(true)))
|
||||
|
|
@ -285,7 +356,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
var capabilities = ClientCapabilities.builder().sampling().build();
|
||||
|
||||
var client = McpClient.async(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.capabilities(capabilities)
|
||||
.sampling(request -> Mono.just(CreateMessageResult.builder().message("test").model("test-model").build()))
|
||||
.build();
|
||||
|
|
@ -309,7 +380,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
Function<CreateMessageRequest, Mono<CreateMessageResult>> samplingHandler = request -> Mono
|
||||
.just(CreateMessageResult.builder().message("test").model("test-model").build());
|
||||
var client = McpClient.async(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.capabilities(capabilities)
|
||||
.sampling(samplingHandler)
|
||||
.build();
|
||||
|
|
@ -326,8 +397,17 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
// Logging Tests
|
||||
// ---------------------------------------
|
||||
|
||||
@Test
|
||||
void testLoggingLevelsWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpAsyncClient.setLoggingLevel(McpSchema.LoggingLevel.DEBUG).block())
|
||||
.isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before setting logging level");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testLoggingLevels() {
|
||||
mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
|
||||
|
||||
// Test all logging levels
|
||||
for (McpSchema.LoggingLevel level : McpSchema.LoggingLevel.values()) {
|
||||
StepVerifier.create(mcpAsyncClient.setLoggingLevel(level)).verifyComplete();
|
||||
|
|
@ -340,7 +420,7 @@ public abstract class AbstractMcpAsyncClientTests {
|
|||
var transport = createMcpTransport();
|
||||
|
||||
var client = McpClient.async(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.loggingConsumer(notification -> Mono.fromRunnable(() -> logReceived.set(true)))
|
||||
.build();
|
||||
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import java.util.Map;
|
|||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
|
||||
import io.modelcontextprotocol.spec.ClientMcpTransport;
|
||||
import io.modelcontextprotocol.spec.McpError;
|
||||
import io.modelcontextprotocol.spec.McpSchema;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
|
||||
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
|
||||
|
|
@ -41,8 +42,6 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
private McpSyncClient mcpSyncClient;
|
||||
|
||||
private static final Duration TIMEOUT = Duration.ofSeconds(10);
|
||||
|
||||
private static final String TEST_MESSAGE = "Hello MCP Spring AI!";
|
||||
|
||||
protected ClientMcpTransport mcpTransport;
|
||||
|
|
@ -53,6 +52,10 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
abstract protected void onClose();
|
||||
|
||||
protected Duration getTimeoutDuration() {
|
||||
return Duration.ofSeconds(2);
|
||||
}
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
onStart();
|
||||
|
|
@ -60,10 +63,9 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
assertThatCode(() -> {
|
||||
mcpSyncClient = McpClient.sync(mcpTransport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.capabilities(ClientCapabilities.builder().roots(true).build())
|
||||
.build();
|
||||
mcpSyncClient.initialize();
|
||||
}).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
|
|
@ -85,8 +87,15 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
.hasMessage("Request timeout must not be null");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListToolsWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpSyncClient.listTools(null)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before listing tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListTools() {
|
||||
mcpSyncClient.initialize();
|
||||
ListToolsResult tools = mcpSyncClient.listTools(null);
|
||||
|
||||
assertThat(tools).isNotNull().satisfies(result -> {
|
||||
|
|
@ -98,8 +107,16 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallToolsWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpSyncClient.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4))))
|
||||
.isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before calling tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallTools() {
|
||||
mcpSyncClient.initialize();
|
||||
CallToolResult toolResult = mcpSyncClient.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4)));
|
||||
|
||||
assertThat(toolResult).isNotNull().satisfies(result -> {
|
||||
|
|
@ -114,13 +131,29 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
});
|
||||
}
|
||||
|
||||
@Test
|
||||
void testPingWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpSyncClient.ping()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before pinging the server");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testPing() {
|
||||
mcpSyncClient.initialize();
|
||||
assertThatCode(() -> mcpSyncClient.ping()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallToolWithoutInitialization() {
|
||||
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE));
|
||||
|
||||
assertThatThrownBy(() -> mcpSyncClient.callTool(callToolRequest)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before calling tools");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testCallTool() {
|
||||
mcpSyncClient.initialize();
|
||||
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE));
|
||||
|
||||
CallToolResult callToolResult = mcpSyncClient.callTool(callToolRequest);
|
||||
|
|
@ -138,13 +171,27 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
assertThatThrownBy(() -> mcpSyncClient.callTool(invalidRequest)).isInstanceOf(Exception.class);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRootsListChangedWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpSyncClient.rootsListChangedNotification()).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before sending roots list changed notification");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRootsListChanged() {
|
||||
mcpSyncClient.initialize();
|
||||
assertThatCode(() -> mcpSyncClient.rootsListChangedNotification()).doesNotThrowAnyException();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResourcesWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpSyncClient.listResources(null)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before listing resources");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResources() {
|
||||
mcpSyncClient.initialize();
|
||||
ListResourcesResult resources = mcpSyncClient.listResources(null);
|
||||
|
||||
assertThat(resources).isNotNull().satisfies(result -> {
|
||||
|
|
@ -168,7 +215,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
var transport = createMcpTransport();
|
||||
|
||||
var client = McpClient.sync(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.roots(new Root("file:///test/path", "test-root"))
|
||||
.build();
|
||||
|
||||
|
|
@ -204,8 +251,17 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
.hasMessageContaining("Root with uri 'nonexistent-uri' not found");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReadResourceWithoutInitialization() {
|
||||
assertThatThrownBy(() -> {
|
||||
Resource resource = new Resource("test://uri", "Test Resource", null, null, null);
|
||||
mcpSyncClient.readResource(resource);
|
||||
}).isInstanceOf(McpError.class).hasMessage("Client must be initialized before reading resources");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReadResource() {
|
||||
mcpSyncClient.initialize();
|
||||
ListResourcesResult resources = mcpSyncClient.listResources(null);
|
||||
|
||||
if (!resources.resources().isEmpty()) {
|
||||
|
|
@ -217,8 +273,15 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResourceTemplatesWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpSyncClient.listResourceTemplates(null)).isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before listing resource templates");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testListResourceTemplates() {
|
||||
mcpSyncClient.initialize();
|
||||
ListResourceTemplatesResult result = mcpSyncClient.listResourceTemplates(null);
|
||||
|
||||
assertThat(result).isNotNull();
|
||||
|
|
@ -250,7 +313,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
|
||||
var transport = createMcpTransport();
|
||||
var client = McpClient.sync(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.toolsChangeConsumer(tools -> toolsNotificationReceived.set(true))
|
||||
.resourcesChangeConsumer(resources -> resourcesNotificationReceived.set(true))
|
||||
.promptsChangeConsumer(prompts -> promptsNotificationReceived.set(true))
|
||||
|
|
@ -266,8 +329,16 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
// Logging Tests
|
||||
// ---------------------------------------
|
||||
|
||||
@Test
|
||||
void testLoggingLevelsWithoutInitialization() {
|
||||
assertThatThrownBy(() -> mcpSyncClient.setLoggingLevel(McpSchema.LoggingLevel.DEBUG))
|
||||
.isInstanceOf(McpError.class)
|
||||
.hasMessage("Client must be initialized before setting logging level");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testLoggingLevels() {
|
||||
mcpSyncClient.initialize();
|
||||
// Test all logging levels
|
||||
for (McpSchema.LoggingLevel level : McpSchema.LoggingLevel.values()) {
|
||||
assertThatCode(() -> mcpSyncClient.setLoggingLevel(level)).doesNotThrowAnyException();
|
||||
|
|
@ -280,7 +351,7 @@ public abstract class AbstractMcpSyncClientTests {
|
|||
var transport = createMcpTransport();
|
||||
|
||||
var client = McpClient.sync(transport)
|
||||
.requestTimeout(TIMEOUT)
|
||||
.requestTimeout(getTimeoutDuration())
|
||||
.loggingConsumer(notification -> logReceived.set(true))
|
||||
.build();
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue