Skip to content

Fix Registry selection, inject Registry into the command #705

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 2 commits into from
May 23, 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 @@ -6,7 +6,6 @@
package software.amazon.smithy.java.mcp.cli;

import java.util.Set;
import software.amazon.smithy.java.mcp.cli.model.Config;
import software.amazon.smithy.java.mcp.cli.model.Location;

/**
Expand All @@ -19,7 +18,8 @@
public abstract class AbstractAddBundle extends SmithyMcpCommand implements ConfigurationCommand {

@Override
public final void execute(Config config) throws Exception {
public final void execute(ExecutionContext context) throws Exception {
var config = context.config();
if (!canOverwrite() && config.getToolBundles().containsKey(getToolBundleName())) {
throw new IllegalArgumentException("Tool bundle " + getToolBundleName()
+ " already exists. Either choose a new name or pass --overwrite to overwrite the existing tool bundle");
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

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

import software.amazon.smithy.java.mcp.cli.model.Config;
import software.amazon.smithy.mcp.bundle.api.Registry;

public record ExecutionContext(Config config, Registry registry) {}
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,15 @@
package software.amazon.smithy.java.mcp.cli;

import java.util.Map;
import java.util.Objects;
import java.util.ServiceLoader;
import java.util.ServiceLoader.Provider;
import java.util.function.Function;
import java.util.stream.Collectors;
import software.amazon.smithy.java.mcp.cli.model.Config;
import software.amazon.smithy.mcp.bundle.api.Registry;

public class RegistryUtils {
class RegistryUtils {

private RegistryUtils() {}

Expand All @@ -25,17 +27,20 @@ private RegistryUtils() {}
.collect(Collectors.toMap(Registry::name, Function.identity()));
}

public static Registry getRegistry(String name) {
static Registry getRegistry(String name, Config config) {

if (name == null) {
return getRegistry();
name = Objects.requireNonNull(config.getDefaultRegistry(),
"Either configure a default registry or the registry name is required.");
}

if (!config.getRegistries().containsKey(name)) {
throw new IllegalArgumentException("The registry '" + name + "' does not exist.");
}

if (JAVA_REGISTRIES.containsKey(name)) {
return JAVA_REGISTRIES.get(name);
}
throw new IllegalStateException("No such registry: " + name);
}

public static Registry getRegistry() {
return JAVA_REGISTRIES.values().iterator().next();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ public abstract class SmithyMcpCommand implements Callable<Integer> {
public final Integer call() throws Exception {
try {
var config = loadOrCreateConfig();
execute(config);
execute(new ExecutionContext(config, RegistryUtils.getRegistry(registryToUse(config), config)));
return 0;
} catch (IllegalArgumentException e) {
System.out.println("Invalid input : [" + e.getMessage() + "]");
Expand All @@ -42,8 +42,12 @@ public final Integer call() throws Exception {
* <p>
* Subclasses must implement this method to provide command-specific functionality.
*
* @param config The MCP configuration
* @param context {@link ExecutionContext}
* @throws Exception If an error occurs during execution
*/
protected abstract void execute(Config config) throws Exception;
protected abstract void execute(ExecutionContext context) throws Exception;

protected String registryToUse(Config config) {
return config.getDefaultRegistry();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
import software.amazon.smithy.java.json.JsonCodec;
import software.amazon.smithy.java.json.JsonSettings;
import software.amazon.smithy.java.mcp.cli.ConfigUtils;
import software.amazon.smithy.java.mcp.cli.RegistryUtils;
import software.amazon.smithy.java.mcp.cli.ExecutionContext;
import software.amazon.smithy.java.mcp.cli.SmithyMcpCommand;
import software.amazon.smithy.java.mcp.cli.model.Config;
import software.amazon.smithy.java.mcp.cli.model.McpServerConfig;
Expand All @@ -35,7 +35,7 @@ public class InstallBundle extends SmithyMcpCommand {

@Option(names = {"-r", "--registry"},
description = "Name of the registry to list the bundles from. If not provided it will use the default registry.")
String registry;
String registryName;

@Option(names = {"-n", "--name"}, description = "Name of the MCP Bundle to install.")
String name;
Expand All @@ -49,11 +49,10 @@ public class InstallBundle extends SmithyMcpCommand {
boolean print;

@Override
protected void execute(Config config) throws IOException {
if (registry != null && !config.getRegistries().containsKey(registry)) {
throw new IllegalArgumentException("The registry '" + registry + "' does not exist.");
}
var bundle = RegistryUtils.getRegistry(registry).getMcpBundle(name);
protected void execute(ExecutionContext context) throws IOException {
var registry = context.registry();
var config = context.config();
var bundle = registry.getMcpBundle(name);
ConfigUtils.addMcpBundle(config, name, bundle);
var newConfig = McpServerConfig.builder().command("mcp-registry").args(List.of("start-server", name)).build();
if (print) {
Expand All @@ -80,4 +79,9 @@ protected void execute(Config config) throws IOException {
StandardOpenOption.TRUNCATE_EXISTING);
}
}

@Override
protected String registryToUse(Config config) {
return registryName;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import static picocli.CommandLine.Command;

import picocli.CommandLine.Option;
import software.amazon.smithy.java.mcp.cli.RegistryUtils;
import software.amazon.smithy.java.mcp.cli.ExecutionContext;
import software.amazon.smithy.java.mcp.cli.SmithyMcpCommand;
import software.amazon.smithy.java.mcp.cli.model.Config;

Expand All @@ -20,12 +20,8 @@ public class ListBundles extends SmithyMcpCommand {
String registryName;

@Override
protected void execute(Config config) {
if (registryName != null && !config.getRegistries().containsKey(registryName)) {
throw new IllegalArgumentException("The registry '" + registryName + "' does not exist.");
}

var registry = registryName != null ? RegistryUtils.getRegistry(registryName) : RegistryUtils.getRegistry();
protected void execute(ExecutionContext context) {
var registry = context.registry();
registry
.listMcpBundles()
.forEach(bundle -> {
Expand All @@ -39,4 +35,9 @@ protected void execute(Config config) {
});

}

@Override
protected String registryToUse(Config config) {
return registryName;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,8 @@
import picocli.CommandLine.Option;
import picocli.CommandLine.Parameters;
import software.amazon.smithy.java.mcp.cli.ConfigUtils;
import software.amazon.smithy.java.mcp.cli.RegistryUtils;
import software.amazon.smithy.java.mcp.cli.ExecutionContext;
import software.amazon.smithy.java.mcp.cli.SmithyMcpCommand;
import software.amazon.smithy.java.mcp.cli.model.Config;
import software.amazon.smithy.java.mcp.cli.model.Location;
import software.amazon.smithy.java.mcp.cli.model.McpBundleConfig;
import software.amazon.smithy.java.mcp.cli.model.SmithyModeledBundleConfig;
Expand All @@ -33,6 +32,7 @@
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;

/**
Expand All @@ -59,30 +59,30 @@ public final class StartServer extends SmithyMcpCommand {
* Loads the requested tool bundles from configuration, creates appropriate services,
* and starts the MCP server.
*
* @param config The MCP configuration
* @param context {@link ExecutionContext}
* @throws IllegalArgumentException If no tool bundles are configured or requested bundles not found
*/
@Override
public void execute(Config config) throws IOException {
public void execute(ExecutionContext context) throws IOException {

var config = context.config();
// By default, load all available tools
if (toolBundles == null || toolBundles.isEmpty()) {
try {
toolBundles = new ArrayList<>(ConfigUtils.loadOrCreateConfig().getToolBundles().keySet());
} catch (IOException e) {
throw new RuntimeException(e);
}
toolBundles = new ArrayList<>(config.getToolBundles().keySet());
}

if (toolBundles.isEmpty() && !registryServer) {
throw new IllegalStateException("No bundles installed");
}

var registry = context.registry();

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

for (var toolBundle : toolBundles) {
var toolBundleConfig = config.getToolBundles().get(toolBundle);
if (toolBundleConfig == null) {
var bundle = RegistryUtils.getRegistry().getMcpBundle(toolBundle);
var bundle = registry.getMcpBundle(toolBundle);
if (bundle == null) {
throw new IllegalArgumentException("Can't find a configured tool bundle for '" + toolBundle + "'.");
} else {
Expand All @@ -107,8 +107,8 @@ public void execute(Config config) throws IOException {

if (registryServer) {
services.add(McpRegistry.builder()
.addInstallServerOperation(new InstallOp())
.addListServersOperation(new ListOp())
.addInstallServerOperation(new InstallOp(registry))
.addListServersOperation(new ListOp(registry))
.build());
}

Expand Down Expand Up @@ -141,9 +141,16 @@ private static Service bundleToService(McpBundleConfig toolBundleConfig) {
}

private static final class ListOp implements ListServersOperation {

private final Registry registry;

private ListOp(Registry registry) {
this.registry = registry;
}

@Override
public ListServersOutput listServers(ListServersInput input, RequestContext context) {
var servers = RegistryUtils.getRegistry()
var servers = registry
.listMcpBundles()
.stream()
.unordered()
Expand All @@ -159,12 +166,19 @@ public ListServersOutput listServers(ListServersInput input, RequestContext cont
}

private final class InstallOp implements InstallServerOperation {

private final Registry registry;

private InstallOp(Registry registry) {
this.registry = registry;
}

@Override
public InstallServerOutput installServer(InstallServerInput input, RequestContext context) {
try {
var config = ConfigUtils.loadOrCreateConfig();
if (!config.getToolBundles().containsKey(input.getServerName())) {
var bundle = RegistryUtils.getRegistry().getMcpBundle(input.getServerName());
var bundle = registry.getMcpBundle(input.getServerName());
if (bundle == null) {
throw new IllegalArgumentException(
"Can't find a configured tool bundle for '" + input.getServerName() + "'.");
Expand Down