MCP-197 Terminate the JVM properly (#138)

Author: nquinquenel

src/main/java/org/sonarsource/sonarqube/mcp/SonarQubeMcpServer.java

@@ -121,7 +121,7 @@ public SonarQubeMcpServer(Map<String, String> environment) {
       this.transportProvider = httpServerManager.getTransportProvider();
     } else {
       this.httpServerManager = null;
-      this.transportProvider = new StdioServerTransportProvider(new ObjectMapper());
+      this.transportProvider = new StdioServerTransportProvider(new ObjectMapper(), this::shutdown);
     }
 
     initializeServices();

src/main/java/org/sonarsource/sonarqube/mcp/transport/StdioServerTransportProvider.java

@@ -39,6 +39,7 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
+import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import reactor.core.publisher.Flux;
@@ -71,30 +72,37 @@ public class StdioServerTransportProvider implements McpServerTransportProvider
 
   private final Sinks.One<Void> inboundReady = Sinks.one();
 
+  private final Runnable shutdownCallback;
+
   /**
    * Creates a new StdioServerTransportProvider with the specified ObjectMapper and
-   * System streams. Will call System.exit(0) when stdin closes (for Docker containers).
+   * System streams. Will call shutdown callback when stdin closes (for Docker containers).
    */
-  public StdioServerTransportProvider(ObjectMapper objectMapper) {
-    this(new JacksonMcpJsonMapper(objectMapper), System.in, System.out);
+  public StdioServerTransportProvider(ObjectMapper objectMapper, Runnable shutdownCallback) {
+    this(new JacksonMcpJsonMapper(objectMapper), System.in, System.out, shutdownCallback);
   }
 
   /**
    * Creates a new StdioServerTransportProvider with the specified ObjectMapper and
-   * streams. Automatically detects if custom streams are used (tests) to disable System.exit().
+   * streams. Automatically detects if custom streams are used (tests) to disable shutdown callback.
    *
    * @param jsonMapper   The JsonMapper to use for JSON serialization/deserialization
    * @param inputStream  The input stream to read from
    * @param outputStream The output stream to write to
    */
   public StdioServerTransportProvider(McpJsonMapper jsonMapper, InputStream inputStream, OutputStream outputStream) {
+    this(jsonMapper, inputStream, outputStream, null);
+  }
+
+  private StdioServerTransportProvider(McpJsonMapper jsonMapper, InputStream inputStream, OutputStream outputStream, @Nullable Runnable shutdownCallback) {
     Assert.notNull(jsonMapper, "The JsonMapper can not be null");
     Assert.notNull(inputStream, "The InputStream can not be null");
     Assert.notNull(outputStream, "The OutputStream can not be null");
 
     this.jsonMapper = jsonMapper;
     this.inputStream = inputStream;
     this.outputStream = outputStream;
+    this.shutdownCallback = shutdownCallback;
   }
 
   @Override
@@ -267,6 +275,16 @@ private void startInboundProcessing() {
               session.close();
             }
             inboundSink.tryEmitComplete();
+            
+            // Trigger graceful shutdown when stdin closes if environment variable is set (avoid stale containers)
+            if (shutdownCallback != null) {
+              logger.info("stdin closed (EOF detected) - initiating graceful shutdown");
+              try {
+                shutdownCallback.run();
+              } catch (Exception e) {
+                logger.error("Error during graceful shutdown", e);
+              }
+            }
           }
         });
       }

src/test/java/org/sonarsource/sonarqube/mcp/SonarQubeMcpServerHttpTest.java

@@ -104,7 +104,7 @@ void should_not_create_http_manager_when_disabled(SonarQubeMcpServerTestHarness
 
     harness.prepareMockWebServer(environment);
 
-    var server = new SonarQubeMcpServer(new StdioServerTransportProvider(new ObjectMapper()), null, environment);
+    var server = new SonarQubeMcpServer(new StdioServerTransportProvider(new ObjectMapper(), null), null, environment);
 
     assertThat(server.getMcpConfiguration().isHttpEnabled()).isFalse();
   }

src/test/java/org/sonarsource/sonarqube/mcp/SonarQubeMcpServerIdeBridgeTest.java

@@ -59,7 +59,7 @@ void should_add_bridge_related_tools_when_ide_bridge_is_available(SonarQubeMcpSe
 
     harness.prepareMockWebServer(environment);
 
-    var server = new SonarQubeMcpServer(new StdioServerTransportProvider(new ObjectMapper()), null, environment);
+    var server = new SonarQubeMcpServer(new StdioServerTransportProvider(new ObjectMapper(), null), null, environment);
     server.start();
 
     var supportedTools = server.getSupportedTools();
@@ -78,7 +78,7 @@ void should_add_analysis_tool_when_ide_bridge_is_not_available(SonarQubeMcpServe
 
     harness.prepareMockWebServer(environment);
 
-    var server = new SonarQubeMcpServer(new StdioServerTransportProvider(new ObjectMapper()), null, environment);
+    var server = new SonarQubeMcpServer(new StdioServerTransportProvider(new ObjectMapper(), null), null, environment);
     server.start();
 
     var supportedTools = server.getSupportedTools();

src/test/java/org/sonarsource/sonarqube/mcp/harness/SonarQubeMcpServerTestHarness.java

@@ -55,7 +55,6 @@ public class SonarQubeMcpServerTestHarness extends TypeBasedParameterResolver<So
     "SONARQUBE_TOKEN", "token");
   private final List<McpSyncClient> clients = new ArrayList<>();
   private Path tempStoragePath;
-  private final List<SonarQubeMcpServer> servers = new ArrayList<>();
   private final MockWebServer mockSonarQubeServer = new MockWebServer();
 
   @Override
@@ -72,15 +71,6 @@ public void beforeEach(ExtensionContext context) {
   public void afterEach(ExtensionContext context) {
     clients.forEach(McpSyncClient::closeGracefully);
     clients.clear();
-    servers.forEach(server -> {
-      try {
-        server.shutdown();
-      } catch (Exception e) {
-        // Log error but continue cleanup
-        System.err.println("Error shutting down server: " + e.getMessage());
-      }
-    });
-    servers.clear();
     cleanupTempStoragePath();
     mockSonarQubeServer.stop();
   }
@@ -142,7 +132,6 @@ public SonarQubeMcpTestClient newClient(Map<String, String> overriddenEnv) {
       .loggingConsumer(SonarQubeMcpServerTestHarness::printLogs).build();
     client.initialize();
     this.clients.add(client);
-    this.servers.add(server);
     return new SonarQubeMcpTestClient(client);
   }
 

src/test/java/org/sonarsource/sonarqube/mcp/transport/StdioServerTransportProviderTest.java

@@ -37,7 +37,7 @@ class StdioServerTransportProviderTest {
    * We need to use reflection to set the session field directly to avoid the background threads.
    */
   private StdioServerTransportProvider createProviderWithMockSession(McpServerSession mockSession) throws Exception {
-    var provider = new StdioServerTransportProvider(new ObjectMapper());
+    var provider = new StdioServerTransportProvider(new ObjectMapper(), null);
     var sessionField = StdioServerTransportProvider.class.getDeclaredField("session");
     sessionField.setAccessible(true);
     sessionField.set(provider, mockSession);
@@ -46,7 +46,7 @@ private StdioServerTransportProvider createProviderWithMockSession(McpServerSess
 
   @Test
   void closeGracefully_should_return_empty_mono_when_session_is_null() {
-    var provider = new StdioServerTransportProvider(new ObjectMapper());
+    var provider = new StdioServerTransportProvider(new ObjectMapper(), null);
 
     var result = provider.closeGracefully();