Skip to content

Commit e31ceab

Browse files
adwsinghrhernandez35
adwsingh
authored and
rhernandez35
committed
Add support for search tools and dynamic tool loading
1 parent 91eed0b commit e31ceab

File tree

16 files changed

+406
-204
lines changed

16 files changed

+406
-204
lines changed

‎examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/BundleMCPServerExample.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ public static void main(String[] args) throws Exception {
1414
var mcpServer = McpServer.builder()
1515
.stdio()
1616
.name("smithy-mcp-server")
17-
.addService(McpBundles.getService(bundle.getValue()))
17+
.addService("dynamodb-mcp", McpBundles.getService(bundle.getValue()))
1818
.build();
1919

2020
mcpServer.start();

‎examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/MCPServerExample.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ public static void main(String[] args) {
1616
var mcpServer = McpServer.builder()
1717
.stdio()
1818
.name("smithy-mcp-server")
19-
.addService(service)
19+
.addService("employee-mcp", service)
2020
.build();
2121

2222
mcpServer.start();

‎examples/mcp-server/src/main/java/software/amazon/smithy/java/example/server/mcp/ProxyMCPExample.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ public static void main(String[] args) {
4141
var mcpServer = McpServer.builder()
4242
.stdio()
4343
.name("smithy-mcp-server")
44-
.addService(mcpService)
44+
.addService("employee-mcp", mcpService)
4545
.build();
4646
mcpServer.start();
4747

‎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-bundle-api/src/main/java/software/amazon/smithy/mcp/bundle/api/SearchableRegistry.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
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+
import java.util.List;
9+
10+
public interface SearchableRegistry extends Registry {
11+
12+
List<RegistryTool> searchTools(String query, int numberOfTools);
13+
}

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

Lines changed: 109 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,14 @@
77

88
import java.io.IOException;
99
import java.util.ArrayList;
10+
import java.util.HashMap;
11+
import java.util.HashSet;
1012
import java.util.List;
1113
import java.util.Map;
14+
import java.util.Set;
1215
import java.util.concurrent.CompletableFuture;
16+
import java.util.concurrent.CopyOnWriteArraySet;
1317
import java.util.function.Supplier;
14-
import java.util.stream.Collectors;
1518
import picocli.CommandLine.Command;
1619
import picocli.CommandLine.Option;
1720
import picocli.CommandLine.Parameters;
@@ -24,22 +27,25 @@
2427
import software.amazon.smithy.java.mcp.cli.model.GenericToolBundleConfig;
2528
import software.amazon.smithy.java.mcp.cli.model.McpBundleConfig;
2629
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;
30+
import software.amazon.smithy.java.mcp.registry.model.InstallToolInput;
31+
import software.amazon.smithy.java.mcp.registry.model.InstallToolOutput;
32+
import software.amazon.smithy.java.mcp.registry.model.SearchToolsInput;
33+
import software.amazon.smithy.java.mcp.registry.model.SearchToolsOutput;
34+
import software.amazon.smithy.java.mcp.registry.model.Tool;
35+
import software.amazon.smithy.java.mcp.registry.service.InstallToolOperation;
3436
import software.amazon.smithy.java.mcp.registry.service.McpRegistry;
37+
import software.amazon.smithy.java.mcp.registry.service.SearchToolsOperation;
3538
import software.amazon.smithy.java.mcp.server.McpServer;
39+
import software.amazon.smithy.java.mcp.server.StdioProxy;
40+
import software.amazon.smithy.java.mcp.server.ToolFilter;
3641
import software.amazon.smithy.java.server.FilteredService;
3742
import software.amazon.smithy.java.server.OperationFilters;
3843
import software.amazon.smithy.java.server.RequestContext;
3944
import software.amazon.smithy.java.server.Service;
4045
import software.amazon.smithy.mcp.bundle.api.McpBundles;
4146
import software.amazon.smithy.mcp.bundle.api.Registry;
42-
import software.amazon.smithy.mcp.bundle.api.model.BundleMetadata;
47+
import software.amazon.smithy.mcp.bundle.api.SearchableRegistry;
48+
import software.amazon.smithy.mcp.bundle.api.model.Bundle;
4349
import software.amazon.smithy.mcp.bundle.api.model.GenericBundle;
4450

4551
/**
@@ -111,7 +117,8 @@ public void execute(ExecutionContext context) throws IOException {
111117
toolBundleConfigs.add(toolBundleConfig);
112118
}
113119

114-
var services = new ArrayList<Service>();
120+
var services = new HashMap<String, Service>();
121+
var allowerServers = new HashSet<String>();
115122
//TODO Till we implement the full MCP spec in MCPServerProxy we can only start a single proxy server.
116123
ProcessIoProxy proxyServer = null;
117124
for (var toolBundleConfig : toolBundleConfigs) {
@@ -120,13 +127,17 @@ public void execute(ExecutionContext context) throws IOException {
120127
if (proxyServer != null) {
121128
throw new IllegalArgumentException("Generic MCP servers cannot be run with other MCP servers");
122129
}
123-
services.add(bundleToService(toolBundleConfig.getValue()));
130+
SmithyModeledBundleConfig smithyConfig = toolBundleConfig.getValue();
131+
services.put(smithyConfig.getName(), bundleToService(smithyConfig));
132+
allowerServers.add(smithyConfig.getName());
124133
}
125134
case genericConfig -> {
126135
if (!services.isEmpty() || proxyServer != null) {
127136
throw new IllegalArgumentException("Generic MCP servers cannot be run with other MCP servers");
128137
}
129138
GenericToolBundleConfig genericToolBundleConfig = toolBundleConfig.getValue();
139+
allowerServers.add(genericToolBundleConfig.getName());
140+
130141
GenericBundle genericBundle =
131142
ConfigUtils.getMcpBundle(genericToolBundleConfig.getName()).getValue();
132143
List<String> combinedArgs = new ArrayList<>();
@@ -151,22 +162,36 @@ public void execute(ExecutionContext context) throws IOException {
151162

152163
ThrowingRunnable awaitCompletion;
153164
Supplier<CompletableFuture<Void>> shutdownMethod;
165+
Set<String> allowedTools = null;
154166
if (proxyServer != null) {
155167
proxyServer.start();
156168
awaitCompletion = proxyServer::awaitCompletion;
157169
shutdownMethod = proxyServer::shutdown;
158170
} else {
159171
if (registryServer) {
160-
services.add(McpRegistry.builder()
161-
.addInstallServerOperation(new InstallOp(registry, config))
162-
.addListServersOperation(new ListOp(registry))
163-
.build());
172+
allowedTools = new CopyOnWriteArraySet<>();
173+
final SearchToolsOperation searchToolsOperation;
174+
if (registry instanceof SearchableRegistry searchableRegistry) {
175+
searchToolsOperation = new SearchOp(searchableRegistry);
176+
allowedTools.add("SearchTools");
177+
} else {
178+
searchToolsOperation = (ignored, ignored1) -> {
179+
throw new UnsupportedOperationException("This registry isn't searchable");
180+
};
181+
}
182+
allowedTools.add("InstallTool");
183+
services.put("registry-mcp",
184+
McpRegistry.builder()
185+
.addInstallToolOperation(new InstallTool(registry, config, allowedTools))
186+
.addSearchToolsOperation(searchToolsOperation)
187+
.build());
164188
}
165189

166190
this.mcpServer =
167191
(McpServer) McpServer.builder()
168192
.stdio()
169-
.addServices(services)
193+
.addService(services)
194+
.toolFilter(getToolFilter(allowerServers, allowedTools))
170195
.name("smithy-mcp-server")
171196
.build();
172197
mcpServer.start();
@@ -186,10 +211,17 @@ public void execute(ExecutionContext context) throws IOException {
186211
}
187212
}
188213

189-
private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
190-
Service service =
191-
McpBundles.getService(ConfigUtils.getMcpBundle(bundleConfig.getName()));
214+
private ToolFilter getToolFilter(Set<String> allowedServers, Set<String> tools) {
215+
return (mcpServerName, toolName) -> {
216+
if (allowedServers.contains(mcpServerName)) {
217+
return true;
218+
}
219+
return tools == null || tools.contains(toolName);
220+
};
221+
}
192222

223+
private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
224+
var service = bundleToService(ConfigUtils.getMcpBundle(bundleConfig.getName()));
193225
if (bundleConfig.hasAllowListedTools() || bundleConfig.hasBlockListedTools()) {
194226
var filter = OperationFilters.allowList(bundleConfig.getAllowListedTools())
195227
.and(OperationFilters.blockList(bundleConfig.getBlockListedTools()));
@@ -198,60 +230,86 @@ private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
198230
return service;
199231
}
200232

201-
private static final class ListOp implements ListServersOperation {
233+
private static Service bundleToService(Bundle bundle) {
234+
return McpBundles.getService(bundle);
235+
}
202236

203-
private final Registry registry;
237+
private static final class SearchOp implements SearchToolsOperation {
238+
239+
private final SearchableRegistry registry;
204240

205-
private ListOp(Registry registry) {
241+
private SearchOp(SearchableRegistry registry) {
206242
this.registry = registry;
207243
}
208244

209245
@Override
210-
public ListServersOutput listServers(ListServersInput input, RequestContext context) {
211-
var servers = registry
212-
.listMcpBundles()
213-
.stream()
214-
.unordered()
215-
.map(Registry.RegistryEntry::getBundleMetadata)
216-
.collect(Collectors.toMap(
217-
BundleMetadata::getName,
218-
bundle -> ServerEntry.builder()
219-
.description(bundle.getDescription())
220-
.build()));
221-
return ListServersOutput.builder()
222-
.servers(servers)
246+
public SearchToolsOutput searchTools(SearchToolsInput input, RequestContext context) {
247+
var tools = registry.searchTools(input.getToolDescription(), input.getNumberOfTools());
248+
return SearchToolsOutput.builder()
249+
.tools(tools.stream()
250+
.map(t -> Tool.builder()
251+
.serverId(t.serverId())
252+
.toolName(t.toolName())
253+
.build())
254+
.toList())
223255
.build();
224256
}
225257
}
226258

227-
private final class InstallOp implements InstallServerOperation {
259+
private final class InstallTool implements InstallToolOperation {
228260

229261
private final Registry registry;
230262
private final Config config;
263+
private final Set<String> installedTools;
231264

232-
private InstallOp(Registry registry, Config config) {
265+
private InstallTool(Registry registry, Config config, Set<String> installedTools) {
233266
this.registry = registry;
234267
this.config = config;
268+
this.installedTools = installedTools;
235269
}
236270

237271
@Override
238-
public InstallServerOutput installServer(InstallServerInput input, RequestContext context) {
239-
try {
240-
if (!config.getToolBundles().containsKey(input.getServerName())) {
241-
var bundle = registry.getMcpBundle(input.getServerName());
242-
if (bundle == null) {
243-
throw new IllegalArgumentException(
244-
"Can't find a configured tool bundle for '" + input.getServerName() + "'.");
245-
} else {
246-
var mcpBundleConfig = ConfigUtils.addMcpBundle(config, input.getServerName(), bundle);
247-
mcpServer.addNewService(bundleToService(mcpBundleConfig.getValue()));
272+
public InstallToolOutput installTool(InstallToolInput input, RequestContext context) {
273+
var tool = input.getTool();
274+
var toolName = tool.getToolName();
275+
var serverId = tool.getServerId();
276+
Bundle bundle;
277+
if (!config.getToolBundles().containsKey(serverId)) {
278+
bundle = registry.getMcpBundle(serverId);
279+
if (bundle == null) {
280+
throw new IllegalArgumentException("Can't find a configured tool bundle for '" + serverId + "'.");
281+
} else {
282+
try {
283+
ConfigUtils.addMcpBundle(config, serverId, bundle);
284+
} catch (Exception e) {
285+
throw new RuntimeException(e);
248286
}
249287
}
250-
} catch (Exception e) {
251-
throw new RuntimeException(e);
288+
} else {
289+
bundle = ConfigUtils.getMcpBundle(serverId);
252290
}
253-
254-
return InstallServerOutput.builder().build();
291+
switch (bundle.type()) {
292+
case genericBundle -> {
293+
GenericBundle genericBundle = bundle.getValue();
294+
if (!mcpServer.containsMcpServer(serverId)) {
295+
mcpServer.addNewProxy(StdioProxy.builder()
296+
.name(serverId)
297+
.command(genericBundle.getRun().getExecutable())
298+
.build());
299+
}
300+
}
301+
case smithyBundle -> {
302+
if (!mcpServer.containsMcpServer(serverId)) {
303+
mcpServer.addNewService(serverId, bundleToService(bundle));
304+
}
305+
}
306+
default -> throw new IllegalArgumentException("Unsupported bundle type: " + bundle.type());
307+
}
308+
mcpServer.refreshTools();
309+
installedTools.add(toolName);
310+
return InstallToolOutput.builder()
311+
.message("Tool " + toolName + " installed. Check your list of tools.")
312+
.build();
255313
}
256314
}
257315

‎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
}

0 commit comments

Comments
 (0)