Fix Registry selection, inject Registry into the command (#705)

* Fix Registry selection, inject Registry into the command

Author: adwsingh

mcp/mcp-cli-api/src/main/java/software/amazon/smithy/java/mcp/cli/AbstractAddBundle.java

@@ -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;
 
 /**
@@ -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");

mcp/mcp-cli-api/src/main/java/software/amazon/smithy/java/mcp/cli/ExecutionContext.java

@@ -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) {}

mcp/mcp-cli-api/src/main/java/software/amazon/smithy/java/mcp/cli/RegistryUtils.java

@@ -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() {}
 
@@ -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();
-    }
 }

mcp/mcp-cli-api/src/main/java/software/amazon/smithy/java/mcp/cli/SmithyMcpCommand.java

@@ -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() + "]");
@@ -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();
+    }
 }

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

@@ -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;
@@ -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;
@@ -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) {
@@ -80,4 +79,9 @@ protected void execute(Config config) throws IOException {
                     StandardOpenOption.TRUNCATE_EXISTING);
         }
     }
+
+    @Override
+    protected String registryToUse(Config config) {
+        return registryName;
+    }
 }

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

@@ -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;
 
@@ -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 -> {
@@ -39,4 +35,9 @@ protected void execute(Config config) {
                 });
 
     }
+
+    @Override
+    protected String registryToUse(Config config) {
+        return registryName;
+    }
 }

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

@@ -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;
@@ -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;
 
 /**
@@ -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 {
@@ -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());
         }
 
@@ -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()
@@ -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() + "'.");