Skip to content

Commit f7ea062

Browse files
chemicLchemicL
committed
Automatic initialization and burst protection
1 parent 122a06a commit f7ea062

File tree

6 files changed

+153
-133
lines changed

6 files changed

+153
-133
lines changed

‎mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientResiliencyTests.java

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,10 @@
11
package io.modelcontextprotocol.client;
22

3-
import com.fasterxml.jackson.databind.ObjectMapper;
43
import eu.rekawek.toxiproxy.Proxy;
54
import eu.rekawek.toxiproxy.ToxiproxyClient;
65
import eu.rekawek.toxiproxy.model.ToxicDirection;
76
import io.modelcontextprotocol.spec.McpClientTransport;
87
import io.modelcontextprotocol.spec.McpSchema;
9-
import org.awaitility.Awaitility;
10-
import org.junit.jupiter.api.AfterAll;
11-
import org.junit.jupiter.api.BeforeAll;
128
import org.junit.jupiter.api.Test;
139
import org.slf4j.Logger;
1410
import org.slf4j.LoggerFactory;
@@ -24,7 +20,6 @@
2420
import java.util.function.Consumer;
2521
import java.util.function.Function;
2622

27-
import static org.assertj.core.api.Assertions.assertThat;
2823
import static org.assertj.core.api.Assertions.assertThatCode;
2924

3025
public abstract class AbstractMcpAsyncClientResiliencyTests {
@@ -69,7 +64,7 @@ public abstract class AbstractMcpAsyncClientResiliencyTests {
6964
host = "http://" + ipAddressViaToxiproxy + ":" + portViaToxiproxy;
7065
}
7166

72-
void disconnect() {
67+
private static void disconnect() {
7368
long start = System.nanoTime();
7469
try {
7570
// disconnect
@@ -86,7 +81,7 @@ void disconnect() {
8681
}
8782
}
8883

89-
void reconnect() {
84+
private static void reconnect() {
9085
long start = System.nanoTime();
9186
try {
9287
proxy.toxics().get("RESET_UPSTREAM").remove();
@@ -100,6 +95,11 @@ void reconnect() {
10095
}
10196
}
10297

98+
private static void restartMcpServer() {
99+
container.stop();
100+
container.start();
101+
}
102+
103103
abstract McpClientTransport createMcpTransport();
104104

105105
protected Duration getRequestTimeout() {
@@ -164,8 +164,7 @@ void testSessionInvalidation() {
164164
withClient(createMcpTransport(), mcpAsyncClient -> {
165165
StepVerifier.create(mcpAsyncClient.initialize()).expectNextCount(1).verifyComplete();
166166

167-
container.stop();
168-
container.start();
167+
restartMcpServer();
169168

170169
// The first try will face the session mismatch exception and the second one
171170
// will go through the re-initialization process.

‎mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpAsyncClientTests.java

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -110,14 +110,16 @@ void tearDown() {
110110
onClose();
111111
}
112112

113-
<T> void verifyInitializationTimeout(Function<McpAsyncClient, Mono<T>> operation, String action) {
113+
<T> void verifyNotificationSucceedsWithImplicitInitialization(Function<McpAsyncClient, Mono<T>> operation,
114+
String action) {
114115
withClient(createMcpTransport(), mcpAsyncClient -> {
115-
StepVerifier.withVirtualTime(() -> operation.apply(mcpAsyncClient))
116-
.expectSubscription()
117-
.thenAwait(getInitializationTimeout())
118-
.consumeErrorWith(e -> assertThat(e).isInstanceOf(McpError.class)
119-
.hasMessage("Client must be initialized before " + action))
120-
.verify();
116+
StepVerifier.create(operation.apply(mcpAsyncClient)).verifyComplete();
117+
});
118+
}
119+
120+
<T> void verifyCallSucceedsWithImplicitInitialization(Function<McpAsyncClient, Mono<T>> operation, String action) {
121+
withClient(createMcpTransport(), mcpAsyncClient -> {
122+
StepVerifier.create(operation.apply(mcpAsyncClient)).expectNextCount(1).verifyComplete();
121123
});
122124
}
123125

@@ -133,7 +135,7 @@ void testConstructorWithInvalidArguments() {
133135

134136
@Test
135137
void testListToolsWithoutInitialization() {
136-
verifyInitializationTimeout(client -> client.listTools(null), "listing tools");
138+
verifyCallSucceedsWithImplicitInitialization(client -> client.listTools(null), "listing tools");
137139
}
138140

139141
@Test
@@ -153,7 +155,7 @@ void testListTools() {
153155

154156
@Test
155157
void testPingWithoutInitialization() {
156-
verifyInitializationTimeout(client -> client.ping(), "pinging the server");
158+
verifyCallSucceedsWithImplicitInitialization(client -> client.ping(), "pinging the server");
157159
}
158160

159161
@Test
@@ -168,7 +170,7 @@ void testPing() {
168170
@Test
169171
void testCallToolWithoutInitialization() {
170172
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", ECHO_TEST_MESSAGE));
171-
verifyInitializationTimeout(client -> client.callTool(callToolRequest), "calling tools");
173+
verifyCallSucceedsWithImplicitInitialization(client -> client.callTool(callToolRequest), "calling tools");
172174
}
173175

174176
@Test
@@ -202,7 +204,7 @@ void testCallToolWithInvalidTool() {
202204

203205
@Test
204206
void testListResourcesWithoutInitialization() {
205-
verifyInitializationTimeout(client -> client.listResources(null), "listing resources");
207+
verifyCallSucceedsWithImplicitInitialization(client -> client.listResources(null), "listing resources");
206208
}
207209

208210
@Test
@@ -233,7 +235,7 @@ void testMcpAsyncClientState() {
233235

234236
@Test
235237
void testListPromptsWithoutInitialization() {
236-
verifyInitializationTimeout(client -> client.listPrompts(null), "listing " + "prompts");
238+
verifyCallSucceedsWithImplicitInitialization(client -> client.listPrompts(null), "listing " + "prompts");
237239
}
238240

239241
@Test
@@ -258,7 +260,7 @@ void testListPrompts() {
258260
@Test
259261
void testGetPromptWithoutInitialization() {
260262
GetPromptRequest request = new GetPromptRequest("simple_prompt", Map.of());
261-
verifyInitializationTimeout(client -> client.getPrompt(request), "getting " + "prompts");
263+
verifyCallSucceedsWithImplicitInitialization(client -> client.getPrompt(request), "getting " + "prompts");
262264
}
263265

264266
@Test
@@ -279,7 +281,7 @@ void testGetPrompt() {
279281

280282
@Test
281283
void testRootsListChangedWithoutInitialization() {
282-
verifyInitializationTimeout(client -> client.rootsListChangedNotification(),
284+
verifyNotificationSucceedsWithImplicitInitialization(client -> client.rootsListChangedNotification(),
283285
"sending roots list changed notification");
284286
}
285287

@@ -354,7 +356,8 @@ void testReadResource() {
354356

355357
@Test
356358
void testListResourceTemplatesWithoutInitialization() {
357-
verifyInitializationTimeout(client -> client.listResourceTemplates(), "listing resource templates");
359+
verifyCallSucceedsWithImplicitInitialization(client -> client.listResourceTemplates(),
360+
"listing resource templates");
358361
}
359362

360363
@Test
@@ -447,8 +450,8 @@ void testInitializeWithAllCapabilities() {
447450

448451
@Test
449452
void testLoggingLevelsWithoutInitialization() {
450-
verifyInitializationTimeout(client -> client.setLoggingLevel(McpSchema.LoggingLevel.DEBUG),
451-
"setting logging level");
453+
verifyNotificationSucceedsWithImplicitInitialization(
454+
client -> client.setLoggingLevel(McpSchema.LoggingLevel.DEBUG), "setting logging level");
452455
}
453456

454457
@Test

‎mcp-test/src/main/java/io/modelcontextprotocol/client/AbstractMcpSyncClientTests.java

Lines changed: 26 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,14 @@
55
package io.modelcontextprotocol.client;
66

77
import java.time.Duration;
8+
import java.util.List;
89
import java.util.Map;
910
import java.util.concurrent.atomic.AtomicBoolean;
1011
import java.util.concurrent.atomic.AtomicReference;
1112
import java.util.function.Consumer;
1213
import java.util.function.Function;
1314

1415
import io.modelcontextprotocol.spec.McpClientTransport;
15-
import io.modelcontextprotocol.spec.McpError;
1616
import io.modelcontextprotocol.spec.McpSchema;
1717
import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
1818
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
@@ -112,33 +112,18 @@ void tearDown() {
112112

113113
static final Object DUMMY_RETURN_VALUE = new Object();
114114

115-
<T> void verifyNotificationTimesOut(Consumer<McpSyncClient> operation, String action) {
116-
verifyCallTimesOut(client -> {
115+
<T> void verifyNotificationSucceedsWithImplicitInitialization(Consumer<McpSyncClient> operation, String action) {
116+
verifyCallSucceedsWithImplicitInitialization(client -> {
117117
operation.accept(client);
118118
return DUMMY_RETURN_VALUE;
119119
}, action);
120120
}
121121

122-
<T> void verifyCallTimesOut(Function<McpSyncClient, T> blockingOperation, String action) {
122+
<T> void verifyCallSucceedsWithImplicitInitialization(Function<McpSyncClient, T> blockingOperation, String action) {
123123
withClient(createMcpTransport(), mcpSyncClient -> {
124-
// This scheduler is not replaced by virtual time scheduler
125-
Scheduler customScheduler = Schedulers.newBoundedElastic(1, 1, "actualBoundedElastic");
126-
127-
StepVerifier.withVirtualTime(() -> Mono.fromSupplier(() -> blockingOperation.apply(mcpSyncClient))
128-
// Offload the blocking call to the real scheduler
129-
.subscribeOn(customScheduler))
130-
.expectSubscription()
131-
// This works without actually waiting but executes all the
132-
// tasks pending execution on the VirtualTimeScheduler.
133-
// It is possible to execute the blocking code from the operation
134-
// because it is blocked on a dedicated Scheduler and the main
135-
// flow is not blocked and uses the VirtualTimeScheduler.
136-
.thenAwait(getInitializationTimeout())
137-
.consumeErrorWith(e -> assertThat(e).isInstanceOf(McpError.class)
138-
.hasMessage("Client must be initialized before " + action))
139-
.verify();
140-
141-
customScheduler.dispose();
124+
StepVerifier.create(Mono.fromSupplier(() -> blockingOperation.apply(mcpSyncClient)))
125+
.expectNextCount(1)
126+
.verifyComplete();
142127
});
143128
}
144129

@@ -154,7 +139,7 @@ void testConstructorWithInvalidArguments() {
154139

155140
@Test
156141
void testListToolsWithoutInitialization() {
157-
verifyCallTimesOut(client -> client.listTools(null), "listing tools");
142+
verifyCallSucceedsWithImplicitInitialization(client -> client.listTools(null), "listing tools");
158143
}
159144

160145
@Test
@@ -175,8 +160,8 @@ void testListTools() {
175160

176161
@Test
177162
void testCallToolsWithoutInitialization() {
178-
verifyCallTimesOut(client -> client.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4))),
179-
"calling tools");
163+
verifyCallSucceedsWithImplicitInitialization(
164+
client -> client.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4))), "calling tools");
180165
}
181166

182167
@Test
@@ -200,7 +185,7 @@ void testCallTools() {
200185

201186
@Test
202187
void testPingWithoutInitialization() {
203-
verifyCallTimesOut(client -> client.ping(), "pinging the server");
188+
verifyCallSucceedsWithImplicitInitialization(client -> client.ping(), "pinging the server");
204189
}
205190

206191
@Test
@@ -214,7 +199,7 @@ void testPing() {
214199
@Test
215200
void testCallToolWithoutInitialization() {
216201
CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE));
217-
verifyCallTimesOut(client -> client.callTool(callToolRequest), "calling tools");
202+
verifyCallSucceedsWithImplicitInitialization(client -> client.callTool(callToolRequest), "calling tools");
218203
}
219204

220205
@Test
@@ -243,7 +228,7 @@ void testCallToolWithInvalidTool() {
243228

244229
@Test
245230
void testRootsListChangedWithoutInitialization() {
246-
verifyNotificationTimesOut(client -> client.rootsListChangedNotification(),
231+
verifyNotificationSucceedsWithImplicitInitialization(client -> client.rootsListChangedNotification(),
247232
"sending roots list changed notification");
248233
}
249234

@@ -257,7 +242,7 @@ void testRootsListChanged() {
257242

258243
@Test
259244
void testListResourcesWithoutInitialization() {
260-
verifyCallTimesOut(client -> client.listResources(null), "listing resources");
245+
verifyCallSucceedsWithImplicitInitialization(client -> client.listResources(null), "listing resources");
261246
}
262247

263248
@Test
@@ -333,8 +318,14 @@ void testRemoveNonExistentRoot() {
333318

334319
@Test
335320
void testReadResourceWithoutInitialization() {
336-
Resource resource = new Resource("test://uri", "Test Resource", null, null, null);
337-
verifyCallTimesOut(client -> client.readResource(resource), "reading resources");
321+
AtomicReference<List<Resource>> resources = new AtomicReference<>();
322+
withClient(createMcpTransport(), mcpSyncClient -> {
323+
mcpSyncClient.initialize();
324+
resources.set(mcpSyncClient.listResources().resources());
325+
});
326+
327+
verifyCallSucceedsWithImplicitInitialization(client -> client.readResource(resources.get().get(0)),
328+
"reading resources");
338329
}
339330

340331
@Test
@@ -355,7 +346,8 @@ void testReadResource() {
355346

356347
@Test
357348
void testListResourceTemplatesWithoutInitialization() {
358-
verifyCallTimesOut(client -> client.listResourceTemplates(null), "listing resource templates");
349+
verifyCallSucceedsWithImplicitInitialization(client -> client.listResourceTemplates(null),
350+
"listing resource templates");
359351
}
360352

361353
@Test
@@ -413,8 +405,8 @@ void testNotificationHandlers() {
413405

414406
@Test
415407
void testLoggingLevelsWithoutInitialization() {
416-
verifyNotificationTimesOut(client -> client.setLoggingLevel(McpSchema.LoggingLevel.DEBUG),
417-
"setting logging level");
408+
verifyNotificationSucceedsWithImplicitInitialization(
409+
client -> client.setLoggingLevel(McpSchema.LoggingLevel.DEBUG), "setting logging level");
418410
}
419411

420412
@Test

0 commit comments

Comments
 (0)