Shut down local MCP server when the client disconnects

Author: rhernandez35

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

@@ -54,7 +54,7 @@ public final Integer call() throws Exception {
     protected String registryToUse(Config config) {
         return config.getDefaultRegistry();
     }
-  
+
     protected final CommandLine.Model.CommandSpec commandSpec() {
         return commandSpec;
     }

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

@@ -8,6 +8,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.stream.Collectors;
 import picocli.CommandLine.Command;
 import picocli.CommandLine.Option;
@@ -52,6 +53,7 @@ public final class StartServer extends SmithyMcpCommand {
     boolean registryServer;
 
     private volatile McpServer mcpServer;
+    private final CountDownLatch shutdownLatch = new CountDownLatch(1);
 
     /**
      * Executes the start-server command.
@@ -115,10 +117,16 @@ public void execute(ExecutionContext context) throws IOException {
         this.mcpServer =
                 (McpServer) McpServer.builder().stdio().addServices(services).name("smithy-mcp-server").build();
         mcpServer.start();
+
+        boolean shutdown = false;
         try {
-            Thread.currentThread().join();
-        } catch (InterruptedException e) {
+            mcpServer.awaitCompletion();
+            shutdown = true;
             mcpServer.shutdown().join();
+        } catch (Exception e) {
+            if (!shutdown) {
+                mcpServer.shutdown().join();
+            }
         }
     }
 

mcp/mcp-server/src/main/java/software/amazon/smithy/java/mcp/server/McpServer.java

@@ -17,6 +17,7 @@
 import java.util.Scanner;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
 import software.amazon.smithy.java.core.schema.Schema;
 import software.amazon.smithy.java.core.schema.SerializableShape;
 import software.amazon.smithy.java.core.schema.SerializableStruct;
@@ -61,6 +62,7 @@ public final class McpServer implements Server {
     private final OutputStream os;
     private final String name;
     private final List<McpServerProxy> proxies;
+    private final CountDownLatch done = new CountDownLatch(1);
 
     McpServer(McpServerBuilder builder) {
         this.tools = createTools(builder.serviceList);
@@ -73,6 +75,8 @@ public final class McpServer implements Server {
                 this.listen();
             } catch (Exception e) {
                 LOG.error("Error handling request", e);
+            } finally {
+                done.countDown();
             }
         });
         listener.setName("stdio-dispatcher");
@@ -350,6 +354,10 @@ public CompletableFuture<Void> shutdown() {
         }
     }
 
+    public void awaitCompletion() throws InterruptedException {
+        done.await();
+    }
+
     private record Tool(ToolInfo toolInfo, Operation operation, McpServerProxy proxy) {
 
         Tool(ToolInfo toolInfo, Operation operation) {