Skip to content

Commit 68e11e8

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

File tree

14 files changed

+322
-163
lines changed

14 files changed

+322
-163
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/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: 95 additions & 49 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;
1011
import java.util.List;
1112
import java.util.Map;
13+
import java.util.Set;
1214
import java.util.concurrent.CompletableFuture;
15+
import java.util.concurrent.CopyOnWriteArraySet;
16+
import java.util.function.Predicate;
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,24 @@
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.model.ToolInfo;
31+
import software.amazon.smithy.java.mcp.registry.model.InstallToolInput;
32+
import software.amazon.smithy.java.mcp.registry.model.InstallToolOutput;
33+
import software.amazon.smithy.java.mcp.registry.model.SearchToolsInput;
34+
import software.amazon.smithy.java.mcp.registry.model.SearchToolsOutput;
35+
import software.amazon.smithy.java.mcp.registry.model.Tool;
36+
import software.amazon.smithy.java.mcp.registry.service.InstallToolOperation;
3437
import software.amazon.smithy.java.mcp.registry.service.McpRegistry;
38+
import software.amazon.smithy.java.mcp.registry.service.SearchToolsOperation;
3539
import software.amazon.smithy.java.mcp.server.McpServer;
40+
import software.amazon.smithy.java.mcp.server.StdioProxy;
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.model.Bundle;
4348
import software.amazon.smithy.mcp.bundle.api.model.GenericBundle;
4449

4550
/**
@@ -111,7 +116,7 @@ public void execute(ExecutionContext context) throws IOException {
111116
toolBundleConfigs.add(toolBundleConfig);
112117
}
113118

114-
var services = new ArrayList<Service>();
119+
var services = new HashMap<String, Service>();
115120
//TODO Till we implement the full MCP spec in MCPServerProxy we can only start a single proxy server.
116121
ProcessIoProxy proxyServer = null;
117122
for (var toolBundleConfig : toolBundleConfigs) {
@@ -120,7 +125,8 @@ public void execute(ExecutionContext context) throws IOException {
120125
if (proxyServer != null) {
121126
throw new IllegalArgumentException("Generic MCP servers cannot be run with other MCP servers");
122127
}
123-
services.add(bundleToService(toolBundleConfig.getValue()));
128+
SmithyModeledBundleConfig smithyConfig = toolBundleConfig.getValue();
129+
services.put(smithyConfig.getName(), bundleToService(smithyConfig));
124130
}
125131
case genericConfig -> {
126132
if (!services.isEmpty() || proxyServer != null) {
@@ -151,20 +157,28 @@ public void execute(ExecutionContext context) throws IOException {
151157

152158
ThrowingRunnable awaitCompletion;
153159
Supplier<CompletableFuture<Void>> shutdownMethod;
160+
Set<String> allowedTools = null;
154161
if (proxyServer != null) {
155162
proxyServer.start();
156163
awaitCompletion = proxyServer::awaitCompletion;
157164
shutdownMethod = proxyServer::shutdown;
158165
} else {
159166
if (registryServer) {
160-
services.add(McpRegistry.builder()
161-
.addInstallServerOperation(new InstallOp(registry, config))
162-
.addListServersOperation(new ListOp(registry))
163-
.build());
167+
allowedTools = new CopyOnWriteArraySet<>();
168+
services.put("registry-mcp",
169+
McpRegistry.builder()
170+
.addInstallToolOperation(new InstallTool(registry, config, allowedTools))
171+
.addSearchToolsOperation(new SearchOp(registry))
172+
.build());
164173
}
165174

166175
this.mcpServer =
167-
(McpServer) McpServer.builder().stdio().addServices(services).name("smithy-mcp-server").build();
176+
(McpServer) McpServer.builder()
177+
.stdio()
178+
.addService(services)
179+
.toolFilter(getToolFilter(allowedTools))
180+
.name("smithy-mcp-server")
181+
.build();
168182
mcpServer.start();
169183
awaitCompletion = mcpServer::awaitCompletion;
170184
shutdownMethod = mcpServer::shutdown;
@@ -182,9 +196,15 @@ public void execute(ExecutionContext context) throws IOException {
182196
}
183197
}
184198

199+
private Predicate<ToolInfo> getToolFilter(Set<String> allowedTools) {
200+
if (allowedTools == null) {
201+
return toolInfo -> true;
202+
}
203+
return toolInfo -> allowedTools.contains(toolInfo.getName());
204+
}
205+
185206
private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
186-
Service service =
187-
McpBundles.getService(ConfigUtils.getMcpBundle(bundleConfig.getName()));
207+
var service = bundleToService(ConfigUtils.getMcpBundle(bundleConfig.getName()));
188208
if (bundleConfig.hasAllowListedTools() || bundleConfig.hasBlockListedTools()) {
189209
var filter = OperationFilters.allowList(bundleConfig.getAllowListedTools())
190210
.and(OperationFilters.blockList(bundleConfig.getBlockListedTools()));
@@ -193,60 +213,86 @@ private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
193213
return service;
194214
}
195215

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

198222
private final Registry registry;
199223

200-
private ListOp(Registry registry) {
224+
private SearchOp(Registry registry) {
201225
this.registry = registry;
202226
}
203227

204228
@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)
229+
public SearchToolsOutput searchTools(SearchToolsInput input, RequestContext context) {
230+
var tools = registry.searchTools(input.getToolDescription(), input.getNumberOfTools());
231+
return SearchToolsOutput.builder()
232+
.tools(tools.stream()
233+
.map(t -> Tool.builder()
234+
.serverId(t.serverId())
235+
.toolName(t.toolName())
236+
.build())
237+
.toList())
218238
.build();
219239
}
220240
}
221241

222-
private final class InstallOp implements InstallServerOperation {
242+
private final class InstallTool implements InstallToolOperation {
223243

224244
private final Registry registry;
225245
private final Config config;
246+
private final Set<String> installedTools;
226247

227-
private InstallOp(Registry registry, Config config) {
248+
private InstallTool(Registry registry, Config config, Set<String> installedTools) {
228249
this.registry = registry;
229250
this.config = config;
251+
this.installedTools = installedTools;
230252
}
231253

232254
@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()));
255+
public InstallToolOutput installTool(InstallToolInput input, RequestContext context) {
256+
var tool = input.getTool();
257+
var toolName = tool.getToolName();
258+
var serverId = tool.getServerId();
259+
Bundle bundle;
260+
if (!config.getToolBundles().containsKey(serverId)) {
261+
bundle = registry.getMcpBundle(serverId);
262+
if (bundle == null) {
263+
throw new IllegalArgumentException("Can't find a configured tool bundle for '" + serverId + "'.");
264+
} else {
265+
try {
266+
ConfigUtils.addMcpBundle(config, serverId, bundle);
267+
} catch (Exception e) {
268+
throw new RuntimeException(e);
243269
}
244270
}
245-
} catch (Exception e) {
246-
throw new RuntimeException(e);
271+
} else {
272+
bundle = ConfigUtils.getMcpBundle(serverId);
247273
}
248-
249-
return InstallServerOutput.builder().build();
274+
installedTools.add(toolName);
275+
switch (bundle.type()) {
276+
case genericBundle -> {
277+
GenericBundle genericBundle = bundle.getValue();
278+
if (!mcpServer.containsMcpServer(serverId)) {
279+
mcpServer.addNewProxy(StdioProxy.builder()
280+
.name(serverId)
281+
.command(genericBundle.getRun().getExecutable())
282+
.build());
283+
}
284+
}
285+
case smithyBundle -> {
286+
if (!mcpServer.containsMcpServer(serverId)) {
287+
mcpServer.addNewService(serverId, bundleToService(bundle));
288+
}
289+
}
290+
default -> throw new IllegalArgumentException("Unsupported bundle type: " + bundle.type());
291+
}
292+
mcpServer.refreshTools();
293+
return InstallToolOutput.builder()
294+
.message("Tool " + toolName + " installed. Check your list of tools.")
295+
.build();
250296
}
251297
}
252298

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