Skip to content

Commit ff96947

Browse files
ultramancodeultramancode
committed
feat(core): implement dynamic tool management with filtering and cursor pagination
- Introduce ToolsRepository for runtime tool management - Support context-aware tool filtering via client exchange - Implement cursor-based pagination for large toolsets - Maintain backward compatibility for existing static registration
1 parent f7a460f commit ff96947

File tree

7 files changed

+442
-90
lines changed

7 files changed

+442
-90
lines changed

‎mcp-core/src/main/java/io/modelcontextprotocol/server/InMemoryToolsRepository.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
/*
2+
* Copyright 2024-2026 the original author or authors.
3+
*/
4+
5+
package io.modelcontextprotocol.server;
6+
7+
import java.util.Comparator;
8+
import java.util.List;
9+
import java.util.concurrent.ConcurrentHashMap;
10+
11+
import io.modelcontextprotocol.spec.McpSchema;
12+
import reactor.core.publisher.Mono;
13+
14+
/**
15+
* Default in-memory implementation of {@link ToolsRepository}.
16+
* <p>
17+
* This implementation stores tools in a thread-safe {@link ConcurrentHashMap}. It
18+
* provides backward compatibility by exposing all registered tools to all clients without
19+
* filtering. Pagination is not supported in this implementation (always returns full
20+
* list), and the cursor parameter is ignored.
21+
* </p>
22+
*/
23+
public class InMemoryToolsRepository implements ToolsRepository {
24+
25+
private final ConcurrentHashMap<String, McpServerFeatures.AsyncToolSpecification> tools = new ConcurrentHashMap<>();
26+
27+
/**
28+
* Create a new empty InMemoryToolsRepository.
29+
*/
30+
public InMemoryToolsRepository() {
31+
}
32+
33+
/**
34+
* Create a new InMemoryToolsRepository initialized with the given tools.
35+
* @param initialTools Collection of tools to register initially
36+
*/
37+
public InMemoryToolsRepository(List<McpServerFeatures.AsyncToolSpecification> initialTools) {
38+
if (initialTools != null) {
39+
for (McpServerFeatures.AsyncToolSpecification tool : initialTools) {
40+
tools.put(tool.tool().name(), tool);
41+
}
42+
}
43+
}
44+
45+
@Override
46+
public Mono<ToolsListResult> listTools(McpAsyncServerExchange exchange, String cursor) {
47+
// Ensure stable tool ordering for MCP clients, as ConcurrentHashMap does not
48+
// guarantee iteration order
49+
List<McpSchema.Tool> toolList = tools.values()
50+
.stream()
51+
.map(McpServerFeatures.AsyncToolSpecification::tool)
52+
.sorted(Comparator.comparing(McpSchema.Tool::name))
53+
.toList();
54+
55+
return Mono.just(new ToolsListResult(toolList, null));
56+
}
57+
58+
@Override
59+
public Mono<McpServerFeatures.AsyncToolSpecification> resolveToolForCall(String name,
60+
McpAsyncServerExchange exchange) {
61+
// Default behavior: finding = allowing.
62+
// Use a custom ToolsRepository implementation for context-aware access control.
63+
return Mono.justOrEmpty(tools.get(name));
64+
}
65+
66+
@Override
67+
public void addTool(McpServerFeatures.AsyncToolSpecification tool) {
68+
// Last-write-wins policy
69+
tools.put(tool.tool().name(), tool);
70+
}
71+
72+
@Override
73+
public void removeTool(String name) {
74+
tools.remove(name);
75+
}
76+
77+
}

‎mcp-core/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java‎

Lines changed: 51 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2024-2024 the original author or authors.
2+
* Copyright 2024-2026 the original author or authors.
33
*/
44

55
package io.modelcontextprotocol.server;
@@ -11,7 +11,6 @@
1111
import java.util.Optional;
1212
import java.util.UUID;
1313
import java.util.concurrent.ConcurrentHashMap;
14-
import java.util.concurrent.CopyOnWriteArrayList;
1514
import java.util.function.BiFunction;
1615

1716
import io.modelcontextprotocol.json.McpJsonMapper;
@@ -103,7 +102,7 @@ public class McpAsyncServer {
103102

104103
private final String instructions;
105104

106-
private final CopyOnWriteArrayList<McpServerFeatures.AsyncToolSpecification> tools = new CopyOnWriteArrayList<>();
105+
private final ToolsRepository toolsRepository;
107106

108107
private final ConcurrentHashMap<String, McpServerFeatures.AsyncResourceSpecification> resources = new ConcurrentHashMap<>();
109108

@@ -136,7 +135,8 @@ public class McpAsyncServer {
136135
this.serverInfo = features.serverInfo();
137136
this.serverCapabilities = features.serverCapabilities().mutate().logging().build();
138137
this.instructions = features.instructions();
139-
this.tools.addAll(withStructuredOutputHandling(jsonSchemaValidator, features.tools()));
138+
this.toolsRepository = initializeToolsRepository(features.toolsRepository(), jsonSchemaValidator,
139+
features.tools());
140140
this.resources.putAll(features.resources());
141141
this.resourceTemplates.putAll(features.resourceTemplates());
142142
this.prompts.putAll(features.prompts());
@@ -153,6 +153,27 @@ public class McpAsyncServer {
153153
requestTimeout, transport, this::asyncInitializeRequestHandler, requestHandlers, notificationHandlers));
154154
}
155155

156+
/**
157+
* Initialize the tools repository, wrapping tools with structured output handling.
158+
*/
159+
private ToolsRepository initializeToolsRepository(ToolsRepository providedRepository,
160+
JsonSchemaValidator jsonSchemaValidator, List<McpServerFeatures.AsyncToolSpecification> initialTools) {
161+
if (providedRepository != null) {
162+
// Add initial tools to the provided repository with structured output
163+
// handling
164+
if (initialTools != null) {
165+
for (McpServerFeatures.AsyncToolSpecification tool : initialTools) {
166+
providedRepository.addTool(withStructuredOutputHandling(jsonSchemaValidator, tool));
167+
}
168+
}
169+
return providedRepository;
170+
}
171+
// Create default in-memory repository with wrapped tools
172+
List<McpServerFeatures.AsyncToolSpecification> wrappedTools = withStructuredOutputHandling(jsonSchemaValidator,
173+
initialTools);
174+
return new InMemoryToolsRepository(wrappedTools);
175+
}
176+
156177
McpAsyncServer(McpStreamableServerTransportProvider mcpTransportProvider, McpJsonMapper jsonMapper,
157178
McpServerFeatures.Async features, Duration requestTimeout,
158179
McpUriTemplateManagerFactory uriTemplateManagerFactory, JsonSchemaValidator jsonSchemaValidator) {
@@ -161,7 +182,8 @@ public class McpAsyncServer {
161182
this.serverInfo = features.serverInfo();
162183
this.serverCapabilities = features.serverCapabilities().mutate().logging().build();
163184
this.instructions = features.instructions();
164-
this.tools.addAll(withStructuredOutputHandling(jsonSchemaValidator, features.tools()));
185+
this.toolsRepository = initializeToolsRepository(features.toolsRepository(), jsonSchemaValidator,
186+
features.tools());
165187
this.resources.putAll(features.resources());
166188
this.resourceTemplates.putAll(features.resourceTemplates());
167189
this.prompts.putAll(features.prompts());
@@ -336,12 +358,7 @@ public Mono<Void> addTool(McpServerFeatures.AsyncToolSpecification toolSpecifica
336358
var wrappedToolSpecification = withStructuredOutputHandling(this.jsonSchemaValidator, toolSpecification);
337359

338360
return Mono.defer(() -> {
339-
// Remove tools with duplicate tool names first
340-
if (this.tools.removeIf(th -> th.tool().name().equals(wrappedToolSpecification.tool().name()))) {
341-
logger.warn("Replace existing Tool with name '{}'", wrappedToolSpecification.tool().name());
342-
}
343-
344-
this.tools.add(wrappedToolSpecification);
361+
this.toolsRepository.addTool(wrappedToolSpecification);
345362
logger.debug("Added tool handler: {}", wrappedToolSpecification.tool().name());
346363

347364
if (this.serverCapabilities.tools().listChanged()) {
@@ -471,7 +488,11 @@ private static McpServerFeatures.AsyncToolSpecification withStructuredOutputHand
471488
* @return A Flux stream of all registered tools
472489
*/
473490
public Flux<Tool> listTools() {
474-
return Flux.fromIterable(this.tools).map(McpServerFeatures.AsyncToolSpecification::tool);
491+
// Note: This method returns all tools without exchange context.
492+
// For context-aware listing, use toolsRepository.listTools(exchange, cursor)
493+
// directly.
494+
return Flux.defer(
495+
() -> toolsRepository.listTools(null, null).flatMapMany(result -> Flux.fromIterable(result.tools())));
475496
}
476497

477498
/**
@@ -488,17 +509,11 @@ public Mono<Void> removeTool(String toolName) {
488509
}
489510

490511
return Mono.defer(() -> {
491-
if (this.tools.removeIf(toolSpecification -> toolSpecification.tool().name().equals(toolName))) {
492-
493-
logger.debug("Removed tool handler: {}", toolName);
494-
if (this.serverCapabilities.tools().listChanged()) {
495-
return notifyToolsListChanged();
496-
}
497-
}
498-
else {
499-
logger.warn("Ignore as a Tool with name '{}' not found", toolName);
512+
this.toolsRepository.removeTool(toolName);
513+
logger.debug("Requested tool removal: {}", toolName);
514+
if (this.serverCapabilities.tools().listChanged()) {
515+
return notifyToolsListChanged();
500516
}
501-
502517
return Mono.empty();
503518
});
504519
}
@@ -513,9 +528,16 @@ public Mono<Void> notifyToolsListChanged() {
513528

514529
private McpRequestHandler<McpSchema.ListToolsResult> toolsListRequestHandler() {
515530
return (exchange, params) -> {
516-
List<Tool> tools = this.tools.stream().map(McpServerFeatures.AsyncToolSpecification::tool).toList();
517-
518-
return Mono.just(new McpSchema.ListToolsResult(tools, null));
531+
// Extract cursor from params if present
532+
String cursor = null;
533+
if (params != null) {
534+
McpSchema.PaginatedRequest paginatedRequest = jsonMapper.convertValue(params,
535+
new TypeRef<McpSchema.PaginatedRequest>() {
536+
});
537+
cursor = paginatedRequest.cursor();
538+
}
539+
return this.toolsRepository.listTools(exchange, cursor)
540+
.map(result -> new McpSchema.ListToolsResult(result.tools(), result.nextCursor()));
519541
};
520542
}
521543

@@ -525,18 +547,10 @@ private McpRequestHandler<CallToolResult> toolsCallRequestHandler() {
525547
new TypeRef<McpSchema.CallToolRequest>() {
526548
});
527549

528-
Optional<McpServerFeatures.AsyncToolSpecification> toolSpecification = this.tools.stream()
529-
.filter(tr -> callToolRequest.name().equals(tr.tool().name()))
530-
.findAny();
531-
532-
if (toolSpecification.isEmpty()) {
533-
return Mono.error(McpError.builder(McpSchema.ErrorCodes.INVALID_PARAMS)
534-
.message("Unknown tool: invalid_tool_name")
535-
.data("Tool not found: " + callToolRequest.name())
536-
.build());
537-
}
538-
539-
return toolSpecification.get().callHandler().apply(exchange, callToolRequest);
550+
return this.toolsRepository.resolveToolForCall(callToolRequest.name(), exchange)
551+
.switchIfEmpty(Mono
552+
.error(McpError.builder(McpSchema.ErrorCodes.METHOD_NOT_FOUND).message("Tool not found").build()))
553+
.flatMap(spec -> spec.callHandler().apply(exchange, callToolRequest));
540554
};
541555
}
542556

0 commit comments

Comments
 (0)