Add support for search tools and dynamic tool loading

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

@@ -14,7 +14,7 @@ public static void main(String[] args) throws Exception {
             var mcpServer = McpServer.builder()
                 .stdio()
                 .name("smithy-mcp-server")
-                .addService(McpBundles.getService(bundle.getValue()))
+                .addService("dynamodb-mcp", McpBundles.getService(bundle.getValue()))
                 .build();
 
             mcpServer.start();

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

@@ -16,7 +16,7 @@ public static void main(String[] args) {
         var mcpServer = McpServer.builder()
                 .stdio()
                 .name("smithy-mcp-server")
-                .addService(service)
+                .addService("employee-mcp", service)
                 .build();
 
         mcpServer.start();

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

@@ -41,7 +41,7 @@ public static void main(String[] args) {
         var mcpServer = McpServer.builder()
                 .stdio()
                 .name("smithy-mcp-server")
-                .addService(mcpService)
+                .addService("employee-mcp", mcpService)
                 .build();
         mcpServer.start();
 

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

@@ -0,0 +1,8 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.mcp.bundle.api;
+
+public record RegistryTool(String serverId, String toolName) {}

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

@@ -0,0 +1,13 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.mcp.bundle.api;
+
+import java.util.List;
+
+public interface SearchableRegistry extends Registry {
+
+    List<RegistryTool> searchTools(String query, int numberOfTools);
+}

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

@@ -7,11 +7,14 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CopyOnWriteArraySet;
 import java.util.function.Supplier;
-import java.util.stream.Collectors;
 import picocli.CommandLine.Command;
 import picocli.CommandLine.Option;
 import picocli.CommandLine.Parameters;
@@ -24,22 +27,25 @@
 import software.amazon.smithy.java.mcp.cli.model.GenericToolBundleConfig;
 import software.amazon.smithy.java.mcp.cli.model.McpBundleConfig;
 import software.amazon.smithy.java.mcp.cli.model.SmithyModeledBundleConfig;
-import software.amazon.smithy.java.mcp.registry.model.InstallServerInput;
-import software.amazon.smithy.java.mcp.registry.model.InstallServerOutput;
-import software.amazon.smithy.java.mcp.registry.model.ListServersInput;
-import software.amazon.smithy.java.mcp.registry.model.ListServersOutput;
-import software.amazon.smithy.java.mcp.registry.model.ServerEntry;
-import software.amazon.smithy.java.mcp.registry.service.InstallServerOperation;
-import software.amazon.smithy.java.mcp.registry.service.ListServersOperation;
+import software.amazon.smithy.java.mcp.registry.model.InstallToolInput;
+import software.amazon.smithy.java.mcp.registry.model.InstallToolOutput;
+import software.amazon.smithy.java.mcp.registry.model.SearchToolsInput;
+import software.amazon.smithy.java.mcp.registry.model.SearchToolsOutput;
+import software.amazon.smithy.java.mcp.registry.model.Tool;
+import software.amazon.smithy.java.mcp.registry.service.InstallToolOperation;
 import software.amazon.smithy.java.mcp.registry.service.McpRegistry;
+import software.amazon.smithy.java.mcp.registry.service.SearchToolsOperation;
 import software.amazon.smithy.java.mcp.server.McpServer;
+import software.amazon.smithy.java.mcp.server.StdioProxy;
+import software.amazon.smithy.java.mcp.server.ToolFilter;
 import software.amazon.smithy.java.server.FilteredService;
 import software.amazon.smithy.java.server.OperationFilters;
 import software.amazon.smithy.java.server.RequestContext;
 import software.amazon.smithy.java.server.Service;
 import software.amazon.smithy.mcp.bundle.api.McpBundles;
 import software.amazon.smithy.mcp.bundle.api.Registry;
-import software.amazon.smithy.mcp.bundle.api.model.BundleMetadata;
+import software.amazon.smithy.mcp.bundle.api.SearchableRegistry;
+import software.amazon.smithy.mcp.bundle.api.model.Bundle;
 import software.amazon.smithy.mcp.bundle.api.model.GenericBundle;
 
 /**
@@ -111,7 +117,8 @@ public void execute(ExecutionContext context) throws IOException {
             toolBundleConfigs.add(toolBundleConfig);
         }
 
-        var services = new ArrayList<Service>();
+        var services = new HashMap<String, Service>();
+        var allowerServers = new HashSet<String>();
         //TODO Till we implement the full MCP spec in MCPServerProxy we can only start a single proxy server.
         ProcessIoProxy proxyServer = null;
         for (var toolBundleConfig : toolBundleConfigs) {
@@ -120,13 +127,17 @@ public void execute(ExecutionContext context) throws IOException {
                     if (proxyServer != null) {
                         throw new IllegalArgumentException("Generic MCP servers cannot be run with other MCP servers");
                     }
-                    services.add(bundleToService(toolBundleConfig.getValue()));
+                    SmithyModeledBundleConfig smithyConfig = toolBundleConfig.getValue();
+                    services.put(smithyConfig.getName(), bundleToService(smithyConfig));
+                    allowerServers.add(smithyConfig.getName());
                 }
                 case genericConfig -> {
                     if (!services.isEmpty() || proxyServer != null) {
                         throw new IllegalArgumentException("Generic MCP servers cannot be run with other MCP servers");
                     }
                     GenericToolBundleConfig genericToolBundleConfig = toolBundleConfig.getValue();
+                    allowerServers.add(genericToolBundleConfig.getName());
+
                     GenericBundle genericBundle =
                             ConfigUtils.getMcpBundle(genericToolBundleConfig.getName()).getValue();
                     List<String> combinedArgs = new ArrayList<>();
@@ -151,22 +162,36 @@ public void execute(ExecutionContext context) throws IOException {
 
         ThrowingRunnable awaitCompletion;
         Supplier<CompletableFuture<Void>> shutdownMethod;
+        Set<String> allowedTools = null;
         if (proxyServer != null) {
             proxyServer.start();
             awaitCompletion = proxyServer::awaitCompletion;
             shutdownMethod = proxyServer::shutdown;
         } else {
             if (registryServer) {
-                services.add(McpRegistry.builder()
-                        .addInstallServerOperation(new InstallOp(registry, config))
-                        .addListServersOperation(new ListOp(registry))
-                        .build());
+                allowedTools = new CopyOnWriteArraySet<>();
+                final SearchToolsOperation searchToolsOperation;
+                if (registry instanceof SearchableRegistry searchableRegistry) {
+                    searchToolsOperation = new SearchOp(searchableRegistry);
+                    allowedTools.add("SearchTools");
+                } else {
+                    searchToolsOperation = (ignored, ignored1) -> {
+                        throw new UnsupportedOperationException("This registry isn't searchable");
+                    };
+                }
+                allowedTools.add("InstallTool");
+                services.put("registry-mcp",
+                        McpRegistry.builder()
+                                .addInstallToolOperation(new InstallTool(registry, config, allowedTools))
+                                .addSearchToolsOperation(searchToolsOperation)
+                                .build());
             }
 
             this.mcpServer =
                     (McpServer) McpServer.builder()
                             .stdio()
-                            .addServices(services)
+                            .addService(services)
+                            .toolFilter(getToolFilter(allowerServers, allowedTools))
                             .name("smithy-mcp-server")
                             .build();
             mcpServer.start();
@@ -186,10 +211,17 @@ public void execute(ExecutionContext context) throws IOException {
         }
     }
 
-    private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
-        Service service =
-                McpBundles.getService(ConfigUtils.getMcpBundle(bundleConfig.getName()));
+    private ToolFilter getToolFilter(Set<String> allowedServers, Set<String> tools) {
+        return (mcpServerName, toolName) -> {
+            if (allowedServers.contains(mcpServerName)) {
+                return true;
+            }
+            return tools == null || tools.contains(toolName);
+        };
+    }
 
+    private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
+        var service = bundleToService(ConfigUtils.getMcpBundle(bundleConfig.getName()));
         if (bundleConfig.hasAllowListedTools() || bundleConfig.hasBlockListedTools()) {
             var filter = OperationFilters.allowList(bundleConfig.getAllowListedTools())
                     .and(OperationFilters.blockList(bundleConfig.getBlockListedTools()));
@@ -198,60 +230,86 @@ private static Service bundleToService(SmithyModeledBundleConfig bundleConfig) {
         return service;
     }
 
-    private static final class ListOp implements ListServersOperation {
+    private static Service bundleToService(Bundle bundle) {
+        return McpBundles.getService(bundle);
+    }
 
-        private final Registry registry;
+    private static final class SearchOp implements SearchToolsOperation {
+
+        private final SearchableRegistry registry;
 
-        private ListOp(Registry registry) {
+        private SearchOp(SearchableRegistry registry) {
             this.registry = registry;
         }
 
         @Override
-        public ListServersOutput listServers(ListServersInput input, RequestContext context) {
-            var servers = registry
-                    .listMcpBundles()
-                    .stream()
-                    .unordered()
-                    .map(Registry.RegistryEntry::getBundleMetadata)
-                    .collect(Collectors.toMap(
-                            BundleMetadata::getName,
-                            bundle -> ServerEntry.builder()
-                                    .description(bundle.getDescription())
-                                    .build()));
-            return ListServersOutput.builder()
-                    .servers(servers)
+        public SearchToolsOutput searchTools(SearchToolsInput input, RequestContext context) {
+            var tools = registry.searchTools(input.getToolDescription(), input.getNumberOfTools());
+            return SearchToolsOutput.builder()
+                    .tools(tools.stream()
+                            .map(t -> Tool.builder()
+                                    .serverId(t.serverId())
+                                    .toolName(t.toolName())
+                                    .build())
+                            .toList())
                     .build();
         }
     }
 
-    private final class InstallOp implements InstallServerOperation {
+    private final class InstallTool implements InstallToolOperation {
 
         private final Registry registry;
         private final Config config;
+        private final Set<String> installedTools;
 
-        private InstallOp(Registry registry, Config config) {
+        private InstallTool(Registry registry, Config config, Set<String> installedTools) {
             this.registry = registry;
             this.config = config;
+            this.installedTools = installedTools;
         }
 
         @Override
-        public InstallServerOutput installServer(InstallServerInput input, RequestContext context) {
-            try {
-                if (!config.getToolBundles().containsKey(input.getServerName())) {
-                    var bundle = registry.getMcpBundle(input.getServerName());
-                    if (bundle == null) {
-                        throw new IllegalArgumentException(
-                                "Can't find a configured tool bundle for '" + input.getServerName() + "'.");
-                    } else {
-                        var mcpBundleConfig = ConfigUtils.addMcpBundle(config, input.getServerName(), bundle);
-                        mcpServer.addNewService(bundleToService(mcpBundleConfig.getValue()));
+        public InstallToolOutput installTool(InstallToolInput input, RequestContext context) {
+            var tool = input.getTool();
+            var toolName = tool.getToolName();
+            var serverId = tool.getServerId();
+            Bundle bundle;
+            if (!config.getToolBundles().containsKey(serverId)) {
+                bundle = registry.getMcpBundle(serverId);
+                if (bundle == null) {
+                    throw new IllegalArgumentException("Can't find a configured tool bundle for '" + serverId + "'.");
+                } else {
+                    try {
+                        ConfigUtils.addMcpBundle(config, serverId, bundle);
+                    } catch (Exception e) {
+                        throw new RuntimeException(e);
                     }
                 }
-            } catch (Exception e) {
-                throw new RuntimeException(e);
+            } else {
+                bundle = ConfigUtils.getMcpBundle(serverId);
             }
-
-            return InstallServerOutput.builder().build();
+            switch (bundle.type()) {
+                case genericBundle -> {
+                    GenericBundle genericBundle = bundle.getValue();
+                    if (!mcpServer.containsMcpServer(serverId)) {
+                        mcpServer.addNewProxy(StdioProxy.builder()
+                                .name(serverId)
+                                .command(genericBundle.getRun().getExecutable())
+                                .build());
+                    }
+                }
+                case smithyBundle -> {
+                    if (!mcpServer.containsMcpServer(serverId)) {
+                        mcpServer.addNewService(serverId, bundleToService(bundle));
+                    }
+                }
+                default -> throw new IllegalArgumentException("Unsupported bundle type: " + bundle.type());
+            }
+            mcpServer.refreshTools();
+            installedTools.add(toolName);
+            return InstallToolOutput.builder()
+                    .message("Tool " + toolName + " installed. Check your list of tools.")
+                    .build();
         }
     }
 

mcp/mcp-schemas/model/main.smithy

@@ -59,6 +59,8 @@ structure InitializeResult with [BaseResult] {
 
     @required
     serverInfo: ServerInfo
+
+    instructions: String
 }
 
 structure Capabilities {

mcp/mcp-schemas/model/registry.smithy

@@ -2,38 +2,65 @@ $version: "2"
 
 namespace smithy.mcp.registry
 
-/// This service provides methods to list and install MCP servers. You can get a list of available servers with
-/// ListServers. Use a server from that method to install a new server with the InstallServer API.
+/// This service provides methods to search MCP Tools and install MCP servers. You can get a list of MCP tools that are most appropriate
+/// for the given task with SearchTools. If the given tool is not already available you can install using the InstallTool api.
+/// 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.
 service McpRegistry {
     operations: [
-        ListServers
-        InstallServer
+        SearchTools
+        InstallTool
     ]
 }
 
-/// List the available MCP servers that you can install
-operation ListServers {
+/// Search MCP Tools that can help to perform a current task or answer a query. This can be invoked multiple times
+operation SearchTools {
+    input := {
+        /// Tool Description based on the dialogue context. Include relevant information like urls, nouns, acronyms etc.
+        /// Example dialogue:
+        /// User: Hi, can you help me create a code review. I use code.amazon.com
+        /// Example Tool Description : "Create a code review on code.amazon.com"
+        toolDescription: String
+
+        /// Number of tools to return based on relevance in descending order of relevance. If not specified, the default is 1
+        @default(1)
+        numberOfTools: Integer
+    }
+
     output := {
-        /// A map of server name to details about that server
+        /// List of MCP tools most relevant for the query, sorted by order of relevance,
+        /// the first tool being the most relevant.
         @required
-        servers: ServerMap
+        tools: Tools
     }
 }
 
-map ServerMap {
-    key: String
-    value: ServerEntry
+list Tools {
+    member: Tool
+}
+
+structure Tool {
+    /// Id of the MCP server this Tool belongs to
+    @required
+    serverId: String
+
+    /// Name of the Tool
+    toolName: String
 }
 
 structure ServerEntry {
     description: String
 }
 
-/// Install a new MCP server for local use.
-operation InstallServer {
+/// Install a new MCP Tool for local use.
+/// 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.
+operation InstallTool {
     input := {
         /// The name of the MCP server to install
         @required
-        serverName: String
+        tool: Tool
+    }
+
+    output := {
+        message: String
     }
 }