Skip to content

Commit 29da18e

Browse files
adwsinghadwsingh
committed
Add support for search tools and dynamic tool loading
1 parent 5a1ebdc commit 29da18e

File tree

10 files changed

+269
-115
lines changed

10 files changed

+269
-115
lines changed

‎mcp/mcp-bundle-api/src/main/java/software/amazon/smithy/mcp/bundle/api/Registry.java

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,10 @@ public BundleMetadata getBundleMetadata() {
4141

4242
List<RegistryEntry> listMcpBundles();
4343

44+
default List<RegistryTool> searchTools(String query, int numberOfTools) {
45+
return List.of();
46+
}
47+
4448
Bundle getMcpBundle(String id);
4549

4650
}

‎mcp/mcp-bundle-api/src/main/java/software/amazon/smithy/mcp/bundle/api/RegistryTool.java

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package software.amazon.smithy.mcp.bundle.api;
7+
8+
public record RegistryTool(String serverId, String toolName) {}

‎mcp/mcp-cli/src/main/java/software/amazon/smithy/java/mcp/cli/commands/StartServer.java

Lines changed: 92 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,17 @@
77

88
import java.io.IOException;
99
import java.util.ArrayList;
10+
import java.util.Arrays;
11+
import java.util.HashMap;
12+
import java.util.HashSet;
1013
import java.util.List;
1114
import java.util.Map;
15+
import java.util.Set;
1216
import java.util.concurrent.CompletableFuture;
17+
import java.util.concurrent.ConcurrentHashMap;
18+
import java.util.concurrent.CopyOnWriteArraySet;
19+
import java.util.function.Predicate;
1320
import java.util.function.Supplier;
14-
import java.util.stream.Collectors;
1521
import picocli.CommandLine.Command;
1622
import picocli.CommandLine.Option;
1723
import picocli.CommandLine.Parameters;
@@ -24,23 +30,26 @@
2430
import software.amazon.smithy.java.mcp.cli.model.GenericToolBundleConfig;
2531
import software.amazon.smithy.java.mcp.cli.model.McpBundleConfig;
2632
import software.amazon.smithy.java.mcp.cli.model.SmithyModeledBundleConfig;
27-
import software.amazon.smithy.java.mcp.registry.model.InstallServerInput;
28-
import software.amazon.smithy.java.mcp.registry.model.InstallServerOutput;
29-
import software.amazon.smithy.java.mcp.registry.model.ListServersInput;
30-
import software.amazon.smithy.java.mcp.registry.model.ListServersOutput;
31-
import software.amazon.smithy.java.mcp.registry.model.ServerEntry;
32-
import software.amazon.smithy.java.mcp.registry.service.InstallServerOperation;
33-
import software.amazon.smithy.java.mcp.registry.service.ListServersOperation;
33+
import software.amazon.smithy.java.mcp.model.ToolInfo;
34+
import software.amazon.smithy.java.mcp.registry.model.InstallToolInput;
35+
import software.amazon.smithy.java.mcp.registry.model.InstallToolOutput;
36+
import software.amazon.smithy.java.mcp.registry.model.SearchToolsInput;
37+
import software.amazon.smithy.java.mcp.registry.model.SearchToolsOutput;
38+
import software.amazon.smithy.java.mcp.registry.model.Tool;
39+
import software.amazon.smithy.java.mcp.registry.service.InstallToolOperation;
3440
import software.amazon.smithy.java.mcp.registry.service.McpRegistry;
41+
import software.amazon.smithy.java.mcp.registry.service.SearchToolsOperation;
3542
import software.amazon.smithy.java.mcp.server.McpServer;
43+
import software.amazon.smithy.java.mcp.server.StdioProxy;
3644
import software.amazon.smithy.java.server.FilteredService;
3745
import software.amazon.smithy.java.server.OperationFilters;
3846
import software.amazon.smithy.java.server.RequestContext;
3947
import software.amazon.smithy.java.server.Service;
4048
import software.amazon.smithy.mcp.bundle.api.McpBundles;
4149
import software.amazon.smithy.mcp.bundle.api.Registry;
42-
import software.amazon.smithy.mcp.bundle.api.model.BundleMetadata;
50+
import software.amazon.smithy.mcp.bundle.api.model.Bundle;
4351
import software.amazon.smithy.mcp.bundle.api.model.GenericBundle;
52+
import software.amazon.smithy.modelbundle.api.model.SmithyBundle;
4453

4554
/**
4655
* Command to start a Smithy MCP server exposing specified tool bundles.
@@ -111,7 +120,7 @@ public void execute(ExecutionContext context) throws IOException {
111120
toolBundleConfigs.add(toolBundleConfig);
112121
}
113122

114-
var services = new ArrayList<Service>();
123+
var services = new HashMap<String, Service>();
115124
//TODO Till we implement the full MCP spec in MCPServerProxy we can only start a single proxy server.
116125
ProcessIoProxy proxyServer = null;
117126
for (var toolBundleConfig : toolBundleConfigs) {
@@ -120,7 +129,8 @@ public void execute(ExecutionContext context) throws IOException {
120129
if (proxyServer != null) {
121130
throw new IllegalArgumentException("Generic MCP servers cannot be run with other MCP servers");
122131
}
123-
services.add(bundleToService(toolBundleConfig.getValue()));
132+
SmithyModeledBundleConfig smithyConfig = toolBundleConfig.getValue();
133+
services.put(smithyConfig.getName(), bundleToService(smithyConfig));
124134
}
125135
case genericConfig -> {
126136
if (!services.isEmpty() || proxyServer != null) {
@@ -151,20 +161,25 @@ public void execute(ExecutionContext context) throws IOException {
151161

152162
ThrowingRunnable awaitCompletion;
153163
Supplier<CompletableFuture<Void>> shutdownMethod;
164+
Set<String> allowedTools = null;
154165
if (proxyServer != null) {
155166
proxyServer.start();
156167
awaitCompletion = proxyServer::awaitCompletion;
157168
shutdownMethod = proxyServer::shutdown;
158169
} else {
159170
if (registryServer) {
160-
services.add(McpRegistry.builder()
161-
.addInstallServerOperation(new InstallOp(registry, config))
162-
.addListServersOperation(new ListOp(registry))
171+
allowedTools = new CopyOnWriteArraySet<>();
172+
services.put("registry-mcp", McpRegistry.builder()
173+
.addInstallToolOperation(new InstallTool(registry, config, allowedTools))
174+
.addSearchToolsOperation(new SearchOp(registry))
163175
.build());
164176
}
165177

166178
this.mcpServer =
167-
(McpServer) McpServer.builder().stdio().addServices(services).name("smithy-mcp-server").build();
179+
(McpServer) McpServer.builder().stdio()
180+
.addService(services)
181+
.toolFilter(getToolFilter(allowedTools))
182+
.name("smithy-mcp-server").build();
168183
mcpServer.start();
169184
awaitCompletion = mcpServer::awaitCompletion;
170185
shutdownMethod = mcpServer::shutdown;
@@ -182,9 +197,15 @@ public void execute(ExecutionContext context) throws IOException {
182197
}
183198
}
184199

200+
private Predicate<ToolInfo> getToolFilter(Set<String> allowedTools) {
201+
if (allowedTools == null) {
202+
return toolInfo -> true;
203+
}
204+
return toolInfo -> allowedTools.contains(toolInfo.getName());
205+
}
206+
185207
private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
186-
Service service =
187-
McpBundles.getService(ConfigUtils.getMcpBundle(bundleConfig.getName()));
208+
var service = bundleToService(ConfigUtils.getMcpBundle(bundleConfig.getName()));
188209
if (bundleConfig.hasAllowListedTools() || bundleConfig.hasBlockListedTools()) {
189210
var filter = OperationFilters.allowList(bundleConfig.getAllowListedTools())
190211
.and(OperationFilters.blockList(bundleConfig.getBlockListedTools()));
@@ -193,60 +214,83 @@ private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
193214
return service;
194215
}
195216

196-
private static final class ListOp implements ListServersOperation {
217+
private static Service bundleToService(Bundle bundle) {
218+
return McpBundles.getService(bundle);
219+
}
220+
221+
private static final class SearchOp implements SearchToolsOperation {
197222

198223
private final Registry registry;
199224

200-
private ListOp(Registry registry) {
225+
private SearchOp(Registry registry) {
201226
this.registry = registry;
202227
}
203228

204229
@Override
205-
public ListServersOutput listServers(ListServersInput input, RequestContext context) {
206-
var servers = registry
207-
.listMcpBundles()
208-
.stream()
209-
.unordered()
210-
.map(Registry.RegistryEntry::getBundleMetadata)
211-
.collect(Collectors.toMap(
212-
BundleMetadata::getName,
213-
bundle -> ServerEntry.builder()
214-
.description(bundle.getDescription())
215-
.build()));
216-
return ListServersOutput.builder()
217-
.servers(servers)
230+
public SearchToolsOutput searchTools(SearchToolsInput input, RequestContext context) {
231+
var tools = registry.searchTools(input.getToolDescription(), input.getNumberOfTools());
232+
return SearchToolsOutput.builder()
233+
.tools(tools.stream()
234+
.map(t -> Tool.builder()
235+
.serverId(t.serverId())
236+
.toolName(t.toolName())
237+
.build())
238+
.toList())
218239
.build();
219240
}
220241
}
221242

222-
private final class InstallOp implements InstallServerOperation {
243+
private final class InstallTool implements InstallToolOperation {
223244

224245
private final Registry registry;
225246
private final Config config;
247+
private final Set<String> installedTools;
226248

227-
private InstallOp(Registry registry, Config config) {
249+
private InstallTool(Registry registry, Config config, Set<String> installedTools) {
228250
this.registry = registry;
229251
this.config = config;
252+
this.installedTools = installedTools;
230253
}
231254

232255
@Override
233-
public InstallServerOutput installServer(InstallServerInput input, RequestContext context) {
234-
try {
235-
if (!config.getToolBundles().containsKey(input.getServerName())) {
236-
var bundle = registry.getMcpBundle(input.getServerName());
237-
if (bundle == null) {
238-
throw new IllegalArgumentException(
239-
"Can't find a configured tool bundle for '" + input.getServerName() + "'.");
240-
} else {
241-
var mcpBundleConfig = ConfigUtils.addMcpBundle(config, input.getServerName(), bundle);
242-
mcpServer.addNewService(bundleToService(mcpBundleConfig.getValue()));
256+
public InstallToolOutput installTool(InstallToolInput input, RequestContext context) {
257+
var tool = input.getTool();
258+
var toolName = tool.getToolName();
259+
var serverId = tool.getServerId();
260+
Bundle bundle;
261+
if (!config.getToolBundles().containsKey(serverId)) {
262+
bundle = registry.getMcpBundle(serverId);
263+
if (bundle == null) {
264+
throw new IllegalArgumentException("Can't find a configured tool bundle for '" + serverId + "'.");
265+
} else {
266+
try {
267+
ConfigUtils.addMcpBundle(config, serverId, bundle);
268+
} catch (Exception e) {
269+
throw new RuntimeException(e);
243270
}
244271
}
245-
} catch (Exception e) {
246-
throw new RuntimeException(e);
272+
} else {
273+
bundle = ConfigUtils.getMcpBundle(serverId);
247274
}
248-
249-
return InstallServerOutput.builder().build();
275+
installedTools.add(toolName);
276+
switch (bundle.type()) {
277+
case genericBundle -> {
278+
GenericBundle genericBundle = bundle.getValue();
279+
if (!mcpServer.containsMcpServer(serverId)) {
280+
mcpServer.addNewProxy(StdioProxy.builder()
281+
.name(serverId)
282+
.command(genericBundle.getRun().getExecutable())
283+
.build());
284+
}
285+
}
286+
case smithyBundle -> {
287+
if (!mcpServer.containsMcpServer(serverId)) {
288+
mcpServer.addNewService(serverId, bundleToService(bundle));
289+
}
290+
}
291+
}
292+
mcpServer.refreshTools();
293+
return InstallToolOutput.builder().message("Tool " + toolName + " installed. Check your list of tools.").build();
250294
}
251295
}
252296

‎mcp/mcp-schemas/model/main.smithy

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,8 @@ structure InitializeResult with [BaseResult] {
5959

6060
@required
6161
serverInfo: ServerInfo
62+
63+
instructions: String
6264
}
6365

6466
structure Capabilities {

‎mcp/mcp-schemas/model/registry.smithy

Lines changed: 41 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,38 +2,65 @@ $version: "2"
22

33
namespace smithy.mcp.registry
44

5-
/// This service provides methods to list and install MCP servers. You can get a list of available servers with
6-
/// ListServers. Use a server from that method to install a new server with the InstallServer API.
5+
/// This service provides methods to search MCP Tools and install MCP servers. You can get a list of MCP tools that are most appropriate
6+
/// for the given task with SearchTools. If the given tool is not already available you can install using the InstallTool api.
7+
/// Be aware that tools installed using InstallTool are available as part of the ToolAssistant MCP server and the MCP serverId returned from search tool needs to be ignored while tool calling.
78
service McpRegistry {
89
operations: [
9-
ListServers
10-
InstallServer
10+
SearchTools
11+
InstallTool
1112
]
1213
}
1314

14-
/// List the available MCP servers that you can install
15-
operation ListServers {
15+
/// Search MCP Tools that can help to perform a current task or answer a query. This can be invoked multiple times
16+
operation SearchTools {
17+
input := {
18+
/// Tool Description based on the dialogue context. Include relevant information like urls, nouns, acronyms etc.
19+
/// Example dialogue:
20+
/// User: Hi, can you help me create a code review. I use code.amazon.com
21+
/// Example Tool Description : "Create a code review on code.amazon.com"
22+
toolDescription: String
23+
24+
/// Number of tools to return based on relevance in descending order of relevance. If not specified, the default is 1
25+
@default(1)
26+
numberOfTools: Integer
27+
}
28+
1629
output := {
17-
/// A map of server name to details about that server
30+
/// List of MCP tools most relevant for the query, sorted by order of relevance,
31+
/// the first tool being the most relevant.
1832
@required
19-
servers: ServerMap
33+
tools: Tools
2034
}
2135
}
2236

23-
map ServerMap {
24-
key: String
25-
value: ServerEntry
37+
list Tools {
38+
member: Tool
39+
}
40+
41+
structure Tool {
42+
/// Id of the MCP server this Tool belongs to
43+
@required
44+
serverId: String
45+
46+
/// Name of the Tool
47+
toolName: String
2648
}
2749

2850
structure ServerEntry {
2951
description: String
3052
}
3153

32-
/// Install a new MCP server for local use.
33-
operation InstallServer {
54+
/// Install a new MCP Tool for local use.
55+
/// Be aware that tools installed using InstallTool are available as part of the ToolAssistant MCP server and the MCP serverId returned from search tool needs to be ignored while tool calling.
56+
operation InstallTool {
3457
input := {
3558
/// The name of the MCP server to install
3659
@required
37-
serverName: String
60+
tool: Tool
61+
}
62+
63+
output := {
64+
message: String
3865
}
3966
}

‎mcp/mcp-schemas/smithy-build.json

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,14 @@
33
"plugins": {
44
"java-type-codegen": {
55
"namespace": "software.amazon.smithy.java.mcp",
6-
"headerFile": "license.txt"
6+
"headerFile": "license.txt",
7+
"runtimeTraits": ["smithy.api#documentation", "smithy.api#examples" ]
78
},
89
"java-server-codegen": {
910
"service": "smithy.mcp.registry#McpRegistry",
1011
"namespace": "software.amazon.smithy.java.mcp.registry",
11-
"headerFile": "license.txt"
12+
"headerFile": "license.txt",
13+
"runtimeTraits": ["smithy.api#documentation", "smithy.api#examples" ]
1214
}
1315
}
1416
}

0 commit comments

Comments
 (0)