Skip to content

Commit 23280a3

Browse files
ultramancodeultramancode
committed
refactor: refine tools repository behavior
1 parent e6afd98 commit 23280a3

File tree

7 files changed

+219
-20
lines changed

7 files changed

+219
-20
lines changed

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
import java.util.concurrent.ConcurrentHashMap;
1010

1111
import io.modelcontextprotocol.spec.McpSchema;
12+
import org.slf4j.Logger;
13+
import org.slf4j.LoggerFactory;
1214
import reactor.core.publisher.Mono;
1315

1416
/**
@@ -22,6 +24,8 @@
2224
*/
2325
public class InMemoryToolsRepository implements ToolsRepository {
2426

27+
private static final Logger logger = LoggerFactory.getLogger(InMemoryToolsRepository.class);
28+
2529
private final ConcurrentHashMap<String, McpServerFeatures.AsyncToolSpecification> tools = new ConcurrentHashMap<>();
2630

2731
/**
@@ -66,12 +70,15 @@ public Mono<McpServerFeatures.AsyncToolSpecification> resolveToolForCall(String
6670
@Override
6771
public void addTool(McpServerFeatures.AsyncToolSpecification tool) {
6872
// Last-write-wins policy
69-
tools.put(tool.tool().name(), tool);
73+
var previous = tools.put(tool.tool().name(), tool);
74+
if (previous != null) {
75+
logger.warn("Replace existing Tool with name '{}'", tool.tool().name());
76+
}
7077
}
7178

7279
@Override
73-
public void removeTool(String name) {
74-
tools.remove(name);
80+
public boolean removeTool(String name) {
81+
return tools.remove(name) != null;
7582
}
7683

7784
}

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

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -509,10 +509,15 @@ public Mono<Void> removeTool(String toolName) {
509509
}
510510

511511
return Mono.defer(() -> {
512-
this.toolsRepository.removeTool(toolName);
513-
logger.debug("Requested tool removal: {}", toolName);
514-
if (this.serverCapabilities.tools().listChanged()) {
515-
return notifyToolsListChanged();
512+
boolean removed = this.toolsRepository.removeTool(toolName);
513+
if (removed) {
514+
logger.debug("Removed tool handler: {}", toolName);
515+
if (this.serverCapabilities.tools().listChanged()) {
516+
return notifyToolsListChanged();
517+
}
518+
}
519+
else {
520+
logger.warn("Ignore as a Tool with name '{}' not found", toolName);
516521
}
517522
return Mono.empty();
518523
});
@@ -548,8 +553,10 @@ private McpRequestHandler<CallToolResult> toolsCallRequestHandler() {
548553
});
549554

550555
return this.toolsRepository.resolveToolForCall(callToolRequest.name(), exchange)
551-
.switchIfEmpty(Mono
552-
.error(McpError.builder(McpSchema.ErrorCodes.METHOD_NOT_FOUND).message("Tool not found").build()))
556+
.switchIfEmpty(Mono.error(McpError.builder(McpSchema.ErrorCodes.INVALID_PARAMS)
557+
.message("Unknown tool: " + callToolRequest.name())
558+
.data("Tool not found: " + callToolRequest.name())
559+
.build()))
553560
.flatMap(spec -> spec.callHandler().apply(exchange, callToolRequest));
554561
};
555562
}

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -241,9 +241,9 @@ private SingleSessionAsyncSpecification(McpServerTransportProvider transportProv
241241
*/
242242
@Override
243243
public McpAsyncServer build() {
244-
var features = new McpServerFeatures.Async(this.serverInfo, this.serverCapabilities, this.tools, null,
245-
this.resources, this.resourceTemplates, this.prompts, this.completions, this.rootsChangeHandlers,
246-
this.instructions);
244+
var features = new McpServerFeatures.Async(this.serverInfo, this.serverCapabilities, this.tools,
245+
this.toolsRepository, this.resources, this.resourceTemplates, this.prompts, this.completions,
246+
this.rootsChangeHandlers, this.instructions);
247247

248248
var jsonSchemaValidator = (this.jsonSchemaValidator != null) ? this.jsonSchemaValidator
249249
: McpJsonDefaults.getSchemaValidator();
@@ -269,9 +269,9 @@ public StreamableServerAsyncSpecification(McpStreamableServerTransportProvider t
269269
*/
270270
@Override
271271
public McpAsyncServer build() {
272-
var features = new McpServerFeatures.Async(this.serverInfo, this.serverCapabilities, this.tools, null,
273-
this.resources, this.resourceTemplates, this.prompts, this.completions, this.rootsChangeHandlers,
274-
this.instructions);
272+
var features = new McpServerFeatures.Async(this.serverInfo, this.serverCapabilities, this.tools,
273+
this.toolsRepository, this.resources, this.resourceTemplates, this.prompts, this.completions,
274+
this.rootsChangeHandlers, this.instructions);
275275
var jsonSchemaValidator = this.jsonSchemaValidator != null ? this.jsonSchemaValidator
276276
: McpJsonDefaults.getSchemaValidator();
277277
return new McpAsyncServer(transportProvider, jsonMapper == null ? McpJsonDefaults.getMapper() : jsonMapper,

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,9 @@ public interface ToolsRepository {
5555
/**
5656
* Remove a tool from the repository at runtime by name.
5757
* @param name The name of the tool to remove
58+
* @return {@code true} if a tool was removed, {@code false} if no matching tool
59+
* existed
5860
*/
59-
void removeTool(String name);
61+
boolean removeTool(String name);
6062

6163
}

‎mcp-core/src/test/java/io/modelcontextprotocol/server/InMemoryToolsRepositoryTests.java‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
/*
2+
* Copyright 2024-2026 the original author or authors.
3+
*/
4+
5+
package io.modelcontextprotocol.server;
6+
7+
import java.util.List;
8+
9+
import ch.qos.logback.classic.Logger;
10+
import ch.qos.logback.classic.spi.ILoggingEvent;
11+
import ch.qos.logback.core.read.ListAppender;
12+
import io.modelcontextprotocol.spec.McpSchema;
13+
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
14+
import org.junit.jupiter.api.AfterEach;
15+
import org.junit.jupiter.api.BeforeEach;
16+
import org.junit.jupiter.api.Test;
17+
import org.slf4j.LoggerFactory;
18+
import reactor.core.publisher.Mono;
19+
20+
import static io.modelcontextprotocol.util.ToolsUtils.EMPTY_JSON_SCHEMA;
21+
import static org.assertj.core.api.Assertions.assertThat;
22+
23+
class InMemoryToolsRepositoryTests {
24+
25+
private final Logger logger = (Logger) LoggerFactory.getLogger(InMemoryToolsRepository.class);
26+
27+
private final ListAppender<ILoggingEvent> logAppender = new ListAppender<>();
28+
29+
@BeforeEach
30+
void setUp() {
31+
logAppender.start();
32+
logger.addAppender(logAppender);
33+
}
34+
35+
@AfterEach
36+
void tearDown() {
37+
logger.detachAppender(logAppender);
38+
logAppender.stop();
39+
}
40+
41+
@Test
42+
void addToolShouldReplaceExistingToolAndLogWarning() {
43+
InMemoryToolsRepository repository = new InMemoryToolsRepository();
44+
45+
repository.addTool(toolSpec("duplicate-tool", "first"));
46+
repository.addTool(toolSpec("duplicate-tool", "second"));
47+
48+
ToolsListResult result = repository.listTools(null, null).block();
49+
assertThat(result).isNotNull();
50+
assertThat(result.tools()).hasSize(1);
51+
assertThat(result.tools().get(0).description()).isEqualTo("second");
52+
assertThat(logAppender.list).hasSize(1);
53+
assertThat(logAppender.list.get(0).getFormattedMessage())
54+
.contains("Replace existing Tool with name 'duplicate-tool'");
55+
}
56+
57+
@Test
58+
void addToolShouldNotLogWarningForFirstRegistration() {
59+
InMemoryToolsRepository repository = new InMemoryToolsRepository();
60+
61+
repository.addTool(toolSpec("new-tool", "desc"));
62+
63+
assertThat(logAppender.list).isEmpty();
64+
}
65+
66+
@Test
67+
void removeToolShouldReturnTrueOnlyWhenToolExists() {
68+
InMemoryToolsRepository repository = new InMemoryToolsRepository();
69+
repository.addTool(toolSpec("remove-me", "desc"));
70+
71+
assertThat(repository.removeTool("remove-me")).isTrue();
72+
assertThat(repository.removeTool("remove-me")).isFalse();
73+
}
74+
75+
private McpServerFeatures.AsyncToolSpecification toolSpec(String name, String description) {
76+
McpSchema.Tool tool = McpSchema.Tool.builder()
77+
.name(name)
78+
.description(description)
79+
.inputSchema(EMPTY_JSON_SCHEMA)
80+
.build();
81+
return McpServerFeatures.AsyncToolSpecification.builder()
82+
.tool(tool)
83+
.callHandler((exchange, request) -> Mono
84+
.just(CallToolResult.builder().content(List.of()).isError(false).build()))
85+
.build();
86+
}
87+
88+
}

‎mcp-test/src/main/java/io/modelcontextprotocol/server/AbstractMcpAsyncServerTests.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -300,8 +300,8 @@ public void addTool(McpServerFeatures.AsyncToolSpecification tool) {
300300
}
301301

302302
@Override
303-
public void removeTool(String name) {
304-
tools.remove(name);
303+
public boolean removeTool(String name) {
304+
return tools.remove(name) != null;
305305
}
306306
};
307307

@@ -351,7 +351,8 @@ public void addTool(McpServerFeatures.AsyncToolSpecification tool) {
351351
}
352352

353353
@Override
354-
public void removeTool(String name) {
354+
public boolean removeTool(String name) {
355+
return false;
355356
}
356357
};
357358

@@ -362,7 +363,7 @@ public void removeTool(String name) {
362363

363364
// First call should have null cursor (first page)
364365
StepVerifier.create(mcpAsyncServer.listTools().collectList()).assertNext(tools -> {
365-
assertThat(receivedCursors).containsExactly(null);
366+
assertThat(receivedCursors).containsExactly((String) null);
366367
}).verifyComplete();
367368

368369
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10))).doesNotThrowAnyException();

‎mcp-test/src/test/java/io/modelcontextprotocol/server/ToolsListChangedNotificationTests.java‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
/*
2+
* Copyright 2024-2026 the original author or authors.
3+
*/
4+
5+
package io.modelcontextprotocol.server;
6+
7+
import java.time.Duration;
8+
import java.util.List;
9+
10+
import io.modelcontextprotocol.MockMcpServerTransport;
11+
import io.modelcontextprotocol.MockMcpServerTransportProvider;
12+
import io.modelcontextprotocol.spec.McpSchema;
13+
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
14+
import io.modelcontextprotocol.spec.McpSchema.ServerCapabilities;
15+
import io.modelcontextprotocol.spec.McpSchema.Tool;
16+
import org.junit.jupiter.api.AfterEach;
17+
import org.junit.jupiter.api.BeforeEach;
18+
import org.junit.jupiter.api.Test;
19+
import reactor.core.publisher.Mono;
20+
import reactor.test.StepVerifier;
21+
22+
import static io.modelcontextprotocol.util.ToolsUtils.EMPTY_JSON_SCHEMA;
23+
import static org.assertj.core.api.Assertions.assertThat;
24+
import static org.assertj.core.api.Assertions.assertThatCode;
25+
26+
class ToolsListChangedNotificationTests {
27+
28+
private MockMcpServerTransport transport;
29+
30+
private MockMcpServerTransportProvider transportProvider;
31+
32+
private McpAsyncServer mcpAsyncServer;
33+
34+
@BeforeEach
35+
void setUp() {
36+
transport = new MockMcpServerTransport();
37+
transportProvider = new MockMcpServerTransportProvider(transport);
38+
}
39+
40+
@AfterEach
41+
void tearDown() {
42+
if (mcpAsyncServer != null) {
43+
assertThatCode(() -> mcpAsyncServer.closeGracefully().block(Duration.ofSeconds(10)))
44+
.doesNotThrowAnyException();
45+
}
46+
}
47+
48+
@Test
49+
void removeNonexistentToolDoesNotNotifyListChanged() {
50+
mcpAsyncServer = McpServer.async(transportProvider)
51+
.serverInfo("test-server", "1.0.0")
52+
.capabilities(ServerCapabilities.builder().tools(true).build())
53+
.build();
54+
55+
transport.clearSentMessages();
56+
57+
StepVerifier.create(mcpAsyncServer.removeTool("missing-tool")).verifyComplete();
58+
59+
assertThat(toolListChangedNotificationCount()).isZero();
60+
}
61+
62+
@Test
63+
void removeExistingToolNotifiesListChanged() {
64+
Tool tool = Tool.builder()
65+
.name("tool1")
66+
.description("tool1 description")
67+
.inputSchema(EMPTY_JSON_SCHEMA)
68+
.build();
69+
70+
mcpAsyncServer = McpServer.async(transportProvider)
71+
.serverInfo("test-server", "1.0.0")
72+
.capabilities(ServerCapabilities.builder().tools(true).build())
73+
.toolCall(tool,
74+
(exchange, request) -> Mono
75+
.just(CallToolResult.builder().content(List.of()).isError(false).build()))
76+
.build();
77+
78+
transport.clearSentMessages();
79+
80+
StepVerifier.create(mcpAsyncServer.removeTool("tool1")).verifyComplete();
81+
82+
assertThat(toolListChangedNotificationCount()).isEqualTo(1);
83+
}
84+
85+
private long toolListChangedNotificationCount() {
86+
return transport.getAllSentMessages()
87+
.stream()
88+
.filter(McpSchema.JSONRPCNotification.class::isInstance)
89+
.map(McpSchema.JSONRPCNotification.class::cast)
90+
.filter(notification -> McpSchema.METHOD_NOTIFICATION_TOOLS_LIST_CHANGED.equals(notification.method()))
91+
.count();
92+
}
93+
94+
}

0 commit comments

Comments
 (0)