Skip to content

Commit e91fe7f

Browse files
authored
Fix: McpAsyncClient#listTools prevent infinite recursion (#631)
* Fix: prevent infinite recursion in listTools() when nextCursor is empty string Signed-off-by: lance <[email protected]>
1 parent 9bbfbe5 commit e91fe7f

File tree

2 files changed

+90
-22
lines changed

2 files changed

+90
-22
lines changed

mcp-core/src/main/java/io/modelcontextprotocol/client/McpAsyncClient.java

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -402,6 +402,7 @@ public Mono<Void> closeGracefully() {
402402
// --------------------------
403403
// Initialization
404404
// --------------------------
405+
405406
/**
406407
* The initialization phase should be the first interaction between client and server.
407408
* The client will ensure it happens in case it has not been explicitly called and in
@@ -448,6 +449,7 @@ public Mono<Object> ping() {
448449
// --------------------------
449450
// Roots
450451
// --------------------------
452+
451453
/**
452454
* Adds a new root to the client's root list.
453455
* @param root The root to add.
@@ -625,13 +627,13 @@ private McpSchema.CallToolResult validateToolResult(String toolName, McpSchema.C
625627
* @return A Mono that emits the list of all tools result
626628
*/
627629
public Mono<McpSchema.ListToolsResult> listTools() {
628-
return this.listTools(McpSchema.FIRST_PAGE)
629-
.expand(result -> (result.nextCursor() != null) ? this.listTools(result.nextCursor()) : Mono.empty())
630-
.reduce(new McpSchema.ListToolsResult(new ArrayList<>(), null), (allToolsResult, result) -> {
631-
allToolsResult.tools().addAll(result.tools());
632-
return allToolsResult;
633-
})
634-
.map(result -> new McpSchema.ListToolsResult(Collections.unmodifiableList(result.tools()), null));
630+
return this.listTools(McpSchema.FIRST_PAGE).expand(result -> {
631+
String next = result.nextCursor();
632+
return (next != null && !next.isEmpty()) ? this.listTools(next) : Mono.empty();
633+
}).reduce(new McpSchema.ListToolsResult(new ArrayList<>(), null), (allToolsResult, result) -> {
634+
allToolsResult.tools().addAll(result.tools());
635+
return allToolsResult;
636+
}).map(result -> new McpSchema.ListToolsResult(Collections.unmodifiableList(result.tools()), null));
635637
}
636638

637639
/**

mcp-core/src/test/java/io/modelcontextprotocol/client/McpAsyncClientTests.java

Lines changed: 81 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -4,24 +4,22 @@
44

55
package io.modelcontextprotocol.client;
66

7+
import java.util.List;
8+
import java.util.Map;
9+
import java.util.Objects;
10+
import java.util.Set;
11+
import java.util.concurrent.atomic.AtomicReference;
12+
import java.util.function.Function;
13+
import java.util.stream.Collectors;
14+
715
import io.modelcontextprotocol.json.TypeRef;
816
import io.modelcontextprotocol.spec.McpClientTransport;
917
import io.modelcontextprotocol.spec.McpSchema;
1018
import io.modelcontextprotocol.spec.ProtocolVersions;
11-
1219
import org.junit.jupiter.api.Test;
13-
14-
import com.fasterxml.jackson.core.JsonProcessingException;
15-
1620
import reactor.core.publisher.Mono;
1721
import reactor.test.StepVerifier;
1822

19-
import java.util.List;
20-
import java.util.Map;
21-
import java.util.Objects;
22-
import java.util.concurrent.atomic.AtomicReference;
23-
import java.util.function.Function;
24-
2523
import static io.modelcontextprotocol.util.McpJsonMapperUtils.JSON_MAPPER;
2624
import static org.assertj.core.api.Assertions.assertThat;
2725
import static org.assertj.core.api.Assertions.assertThatCode;
@@ -40,8 +38,7 @@ class McpAsyncClientTests {
4038

4139
private static final String CONTEXT_KEY = "context.key";
4240

43-
private McpClientTransport createMockTransportForToolValidation(boolean hasOutputSchema, boolean invalidOutput)
44-
throws JsonProcessingException {
41+
private McpClientTransport createMockTransportForToolValidation(boolean hasOutputSchema, boolean invalidOutput) {
4542

4643
// Create tool with or without output schema
4744
Map<String, Object> inputSchemaMap = Map.of("type", "object", "properties",
@@ -182,7 +179,7 @@ public java.lang.reflect.Type getType() {
182179
}
183180

184181
@Test
185-
void testCallToolWithOutputSchemaValidationSuccess() throws JsonProcessingException {
182+
void testCallToolWithOutputSchemaValidationSuccess() {
186183
McpClientTransport transport = createMockTransportForToolValidation(true, false);
187184

188185
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
@@ -204,7 +201,7 @@ void testCallToolWithOutputSchemaValidationSuccess() throws JsonProcessingExcept
204201
}
205202

206203
@Test
207-
void testCallToolWithNoOutputSchemaSuccess() throws JsonProcessingException {
204+
void testCallToolWithNoOutputSchemaSuccess() {
208205
McpClientTransport transport = createMockTransportForToolValidation(false, false);
209206

210207
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
@@ -226,7 +223,7 @@ void testCallToolWithNoOutputSchemaSuccess() throws JsonProcessingException {
226223
}
227224

228225
@Test
229-
void testCallToolWithOutputSchemaValidationFailure() throws JsonProcessingException {
226+
void testCallToolWithOutputSchemaValidationFailure() {
230227
McpClientTransport transport = createMockTransportForToolValidation(true, true);
231228

232229
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
@@ -241,4 +238,73 @@ void testCallToolWithOutputSchemaValidationFailure() throws JsonProcessingExcept
241238
StepVerifier.create(client.closeGracefully()).verifyComplete();
242239
}
243240

241+
@Test
242+
void testListToolsWithEmptyCursor() {
243+
McpSchema.Tool addTool = McpSchema.Tool.builder().name("add").description("calculate add").build();
244+
McpSchema.Tool subtractTool = McpSchema.Tool.builder()
245+
.name("subtract")
246+
.description("calculate subtract")
247+
.build();
248+
McpSchema.ListToolsResult mockToolsResult = new McpSchema.ListToolsResult(List.of(addTool, subtractTool), "");
249+
250+
McpClientTransport transport = new McpClientTransport() {
251+
Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler;
252+
253+
@Override
254+
public Mono<Void> connect(
255+
Function<Mono<McpSchema.JSONRPCMessage>, Mono<McpSchema.JSONRPCMessage>> handler) {
256+
return Mono.deferContextual(ctx -> {
257+
this.handler = handler;
258+
return Mono.empty();
259+
});
260+
}
261+
262+
@Override
263+
public Mono<Void> closeGracefully() {
264+
return Mono.empty();
265+
}
266+
267+
@Override
268+
public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
269+
if (!(message instanceof McpSchema.JSONRPCRequest request)) {
270+
return Mono.empty();
271+
}
272+
273+
McpSchema.JSONRPCResponse response;
274+
if (McpSchema.METHOD_INITIALIZE.equals(request.method())) {
275+
response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), MOCK_INIT_RESULT,
276+
null);
277+
}
278+
else if (McpSchema.METHOD_TOOLS_LIST.equals(request.method())) {
279+
response = new McpSchema.JSONRPCResponse(McpSchema.JSONRPC_VERSION, request.id(), mockToolsResult,
280+
null);
281+
}
282+
else {
283+
return Mono.empty();
284+
}
285+
286+
return handler.apply(Mono.just(response)).then();
287+
}
288+
289+
@Override
290+
public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
291+
return JSON_MAPPER.convertValue(data, new TypeRef<>() {
292+
@Override
293+
public java.lang.reflect.Type getType() {
294+
return typeRef.getType();
295+
}
296+
});
297+
}
298+
};
299+
300+
McpAsyncClient client = McpClient.async(transport).enableCallToolSchemaCaching(true).build();
301+
302+
Mono<McpSchema.ListToolsResult> mono = client.listTools();
303+
McpSchema.ListToolsResult toolsResult = mono.block();
304+
assertThat(toolsResult).isNotNull();
305+
306+
Set<String> names = toolsResult.tools().stream().map(McpSchema.Tool::name).collect(Collectors.toSet());
307+
assertThat(names).containsExactlyInAnyOrder("subtract", "add");
308+
}
309+
244310
}

0 commit comments

Comments
 (0)