Skip to content

Rename mcp-registry server to tool-assistant #831

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 30, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,11 @@

package software.amazon.smithy.java.mcp.cli.commands;

import static picocli.CommandLine.*;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

star import

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

eh, I guess it's fine if checkstyle didn't reject it


import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
Expand All @@ -15,8 +18,8 @@
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;
import picocli.CommandLine.Unmatched;
import software.amazon.smithy.java.mcp.cli.ConfigUtils;
Expand All @@ -27,17 +30,17 @@
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.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.mcp.toolassistant.model.InstallToolInput;
import software.amazon.smithy.java.mcp.toolassistant.model.InstallToolOutput;
import software.amazon.smithy.java.mcp.toolassistant.model.SearchToolsInput;
import software.amazon.smithy.java.mcp.toolassistant.model.SearchToolsOutput;
import software.amazon.smithy.java.mcp.toolassistant.model.Tool;
import software.amazon.smithy.java.mcp.toolassistant.service.InstallToolOperation;
import software.amazon.smithy.java.mcp.toolassistant.service.SearchToolsOperation;
import software.amazon.smithy.java.mcp.toolassistant.service.ToolAssistant;
import software.amazon.smithy.java.server.FilteredService;
import software.amazon.smithy.java.server.OperationFilters;
import software.amazon.smithy.java.server.RequestContext;
Expand All @@ -58,11 +61,12 @@
@Command(name = "start-server", description = "Starts an MCP server.")
public final class StartServer extends SmithyMcpCommand {

@Parameters(paramLabel = "TOOL_BUNDLES", description = "Name(s) of the Tool Bundles to expose in this MCP Server.")
List<String> toolBundles;
@Parameters(paramLabel = "MCP_SERVER_IDS",
description = "Id(s) of MCP Servers to start and expose as a single MCP server.")
List<String> mcpServerIds = List.of();

@Option(names = "--registry-server", description = "Serve the registry as an MCP server")
boolean registryServer;
@Option(names = {"--tool-assistant", "-ts"}, description = "Exposes a Tool Assistant to Search and Install tools")
boolean toolAssistant;

@Unmatched
List<String> additionalArgs;
Expand All @@ -82,9 +86,9 @@ public final class StartServer extends SmithyMcpCommand {
public void execute(ExecutionContext context) throws IOException {

var config = context.config();
// By default, load all available tools
if (toolBundles == null || toolBundles.isEmpty()) {
toolBundles = config.getToolBundles()
// By default, load all available tools only if not in tool-assistant mode.
if (!toolAssistant && mcpServerIds.isEmpty()) {
mcpServerIds = config.getToolBundles()
.entrySet()
.stream()
.filter(entry -> {
Expand All @@ -96,22 +100,22 @@ public void execute(ExecutionContext context) throws IOException {
.toList();
}

if (toolBundles.isEmpty() && !registryServer) {
throw new IllegalArgumentException("No bundles installed");
if (!toolAssistant && mcpServerIds.isEmpty()) {
throw new IllegalArgumentException("No MCP servers installed");
}

var registry = context.registry();

List<McpBundleConfig> toolBundleConfigs = new ArrayList<>(toolBundles.size());
List<McpBundleConfig> toolBundleConfigs = new ArrayList<>(mcpServerIds.size());

for (var toolBundle : toolBundles) {
var toolBundleConfig = config.getToolBundles().get(toolBundle);
for (var id : mcpServerIds) {
var toolBundleConfig = config.getToolBundles().get(id);
if (toolBundleConfig == null) {
var bundle = registry.getMcpBundle(toolBundle);
var bundle = registry.getMcpBundle(id);
if (bundle == null) {
throw new IllegalArgumentException("Can't find a configured tool bundle for '" + toolBundle + "'.");
throw new IllegalArgumentException("Can't find a configured MCP server for '" + id + "'.");
} else {
toolBundleConfig = ConfigUtils.addMcpBundle(config, toolBundle, bundle);
toolBundleConfig = ConfigUtils.addMcpBundle(config, id, bundle);
}
}
toolBundleConfigs.add(toolBundleConfig);
Expand Down Expand Up @@ -168,7 +172,7 @@ public void execute(ExecutionContext context) throws IOException {
awaitCompletion = proxyServer::awaitCompletion;
shutdownMethod = proxyServer::shutdown;
} else {
if (registryServer) {
if (toolAssistant) {
allowedTools = new CopyOnWriteArraySet<>();
final SearchToolsOperation searchToolsOperation;
if (registry instanceof SearchableRegistry searchableRegistry) {
Expand All @@ -180,8 +184,8 @@ public void execute(ExecutionContext context) throws IOException {
};
}
allowedTools.add("InstallTool");
services.put("registry-mcp",
McpRegistry.builder()
services.put("tool-assistant",
ToolAssistant.builder()
.addInstallToolOperation(new InstallTool(registry, config, allowedTools))
.addSearchToolsOperation(searchToolsOperation)
.build());
Expand Down Expand Up @@ -248,8 +252,7 @@ public SearchToolsOutput searchTools(SearchToolsInput input, RequestContext cont
return SearchToolsOutput.builder()
.tools(tools.stream()
.map(t -> Tool.builder()
.serverId(t.serverId())
.toolName(t.toolName())
.toolName(t.serverId() + "__" + t.toolName())
.build())
.toList())
.build();
Expand All @@ -270,9 +273,13 @@ private InstallTool(Registry registry, Config config, Set<String> installedTools

@Override
public InstallToolOutput installTool(InstallToolInput input, RequestContext context) {
var tool = input.getTool();
var toolName = tool.getToolName();
var serverId = tool.getServerId();
var tool = input.getToolName();
var toolParts = tool.split("__");
if (toolParts.length < 2) {
throw new IllegalArgumentException("Invalid tool name");
}
var serverId = toolParts[0];
var toolName = Arrays.stream(toolParts).skip(1).collect(Collectors.joining("__"));
Bundle bundle;
if (!config.getToolBundles().containsKey(serverId)) {
bundle = registry.getMcpBundle(serverId);
Expand Down
66 changes: 0 additions & 66 deletions mcp/mcp-schemas/model/registry.smithy

This file was deleted.

74 changes: 74 additions & 0 deletions mcp/mcp-schemas/model/tool-assistant.smithy
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
$version: "2"

namespace smithy.mcp.toolassistant

/// This service provides methods to search MCP Tools and install MCP servers.
/// 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 ToolAssistant {
operations: [
SearchTools
InstallTool
]
}

/// You MUST use the SearchTools MCP tool whenever the current toolset is not sufficient to complete the task.
/// Do not attempt to proceed with incomplete or inadequate tools.
/// SearchTools returns a ranked list of tools relevant to the task description—always use it to expand your capabilities.
/// If any required tool is missing, immediately install it using the InstallTool API.
/// This process can and should be repeated: keep invoking SearchTools and InstallTool until you have the tools needed to complete the task successfully.
/// Do not rely solely on the preloaded tools. Tool discovery and installation is a core part of your workflow.
operation SearchTools {
input := {
/// Generate a concise tool description based on the current dialogue context.
///
/// Focus on the user’s intent, including specific actions, platforms, or systems mentioned.
///
/// Include concrete nouns, URLs, acronyms, or product names that help identify the task.
///
/// The goal is to create a search query that helps SearchTools return the most relevant tools.
///
/// Avoid vague or generic phrasing—be as specific and task-oriented as possible.
///
// Example
///
/// Dialogue: "Hi, can you help me create a code review? I use code.amazon.com"
///
/// 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 := {
/// List of MCP tools most relevant for the query, sorted by order of relevance,
/// the first tool being the most relevant.
@required
tools: Tools
}
}

list Tools {
member: Tool
}

structure Tool {
/// Name of this tool
@required
toolName: String
}

/// Install a new MCP Tool for local use.
/// Be aware that tools installed using InstallTool are available as part of the ToolAssistant MCP server.
operation InstallTool {
input := {
/// The name of the MCP tool to install
@required
toolName: String
}

output := {
message: String
}
}
4 changes: 2 additions & 2 deletions mcp/mcp-schemas/smithy-build.json
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@
"runtimeTraits": ["smithy.api#documentation", "smithy.api#examples" ]
},
"java-server-codegen": {
"service": "smithy.mcp.registry#McpRegistry",
"namespace": "software.amazon.smithy.java.mcp.registry",
"service": "smithy.mcp.toolassistant#ToolAssistant",
"namespace": "software.amazon.smithy.java.mcp.toolassistant",
"headerFile": "license.txt",
"runtimeTraits": ["smithy.api#documentation", "smithy.api#examples" ]
}
Expand Down
Loading