refactor: improve process handling and graceful shutdown

- Enhance process handling in StdioServerTransport with proper cleanup
- Refactor graceful shutdown across McpSession implementations
- Add proper JSON serialization annotations for Role and LoggingLevel enums
- Fix typo in setInboundMessageHandler method name
- Improve test coverage using StepVerifier
- Add documentation and cleanup code

Author: chemicL

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/McpSyncClient.java

@@ -44,7 +44,7 @@ public McpSyncClient(McpAsyncClient delegate) {
 
 	@Override
 	public void close() {
-		this.delegate.closeGracefully(Duration.ofMillis(DEFAULT_CLOSE_TIMEOUT_MS)).block();
+		this.delegate.closeGracefully().block(Duration.ofMillis(DEFAULT_CLOSE_TIMEOUT_MS));
 	}
 
 	/**

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/client/stdio/StdioServerTransport.java

@@ -22,6 +22,7 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executors;
 
 import com.fasterxml.jackson.core.type.TypeReference;
@@ -87,6 +88,11 @@ public StdioServerTransport(ServerParameters params, ObjectMapper objectMapper)
 
 	@Override
 	public void start() {
+		// Let's kick off the abstraction layer that will consume the logical messages
+		// pushed via sinks. The code that follows is actually feeding the sinks with
+		// data.
+		super.start();
+
 		// Prepare command and environment
 		List<String> fullCommand = new ArrayList<>();
 		fullCommand.add(params.getCommand());
@@ -234,39 +240,43 @@ public Mono<Void> closeGracefully() {
 
 		this.isRunning = false;
 
-		return Mono.whenDelayError(Mono.fromRunnable(() -> {
+		return Mono.fromFuture(() -> {
+			System.out.println("Sending TERM to process");
+			if (this.process != null) {
+				this.process.destroy();
+				return process.onExit();
+			}
+			else {
+				return CompletableFuture.failedFuture(new RuntimeException("Process not started"));
+			}
+		}).doOnNext(process -> {
+			if (process.exitValue() != 0) {
+				System.out.println("Process terminated with code " + process.exitValue());
+			}
+		}).then(Mono.whenDelayError(Mono.fromRunnable(() -> {
 			try {
 				this.processErrorReader.close();
 			}
 			catch (IOException e) {
 				throw new RuntimeException(e);
 			}
-		}).subscribeOn(errorScheduler), Mono.fromRunnable(() -> {
+		}), Mono.fromRunnable(() -> {
 			try {
 				this.processReader.close();
 			}
 			catch (IOException e) {
 				throw new RuntimeException(e);
 			}
-		}).subscribeOn(inboundScheduler), Mono.fromRunnable(() -> {
+		}), Mono.fromRunnable(() -> {
 			try {
 				this.processWriter.close();
 			}
 			catch (IOException e) {
 				throw new RuntimeException(e);
 			}
-		}).subscribeOn(outboundScheduler))
+		})))
 			.then(Mono.whenDelayError(inboundScheduler.disposeGracefully(), outboundScheduler.disposeGracefully(),
 					errorScheduler.disposeGracefully()))
-			.then(Mono.fromRunnable(() -> {
-				if (this.process != null) {
-					var process = this.process.destroyForcibly();
-					if (process.exitValue() != 0) {
-						throw new RuntimeException("Failed to destroy process");
-					}
-				}
-			}))
-			.then()
 			.subscribeOn(Schedulers.boundedElastic());
 	}
 

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/AbstractMcpTransport.java

@@ -29,6 +29,9 @@
  * @author Christian Tzolov
  * @author Dariusz Jędrzejczyk
  */
+// TODO: The API is not finalized. Implementing SSE will dictate a more concrete
+// structure and relationship between the base and the actual implementation including
+// proper encapsulation and state management.
 public abstract class AbstractMcpTransport implements McpTransport {
 
 	protected final ObjectMapper objectMapper;
@@ -53,17 +56,9 @@ public AbstractMcpTransport(ObjectMapper objectMapper) {
 		Assert.notNull(objectMapper, "ObjectMapper must not be null");
 		this.objectMapper = objectMapper;
 
-		// TODO: consider the effects of buffering here -> the inter-process pipes are
-		// independent and the notifications can flood the client/server.
-		// Potentially, the interest in reading could be communicated from one party
-		// to the other so the Blocking IO Threads can pause consuming the stream
-		// buffers when there is no expectation for reading.
-
 		this.errorSink = Sinks.many().unicast().onBackpressureBuffer();
 		this.inboundSink = Sinks.many().unicast().onBackpressureBuffer();
 		this.outboundSink = Sinks.many().unicast().onBackpressureBuffer();
-
-		this.handleIncomingMessages();
 	}
 
 	public ObjectMapper getObjectMapper() {
@@ -77,7 +72,6 @@ private void handleIncomingMessages() {
 	private void handleIncomingErrors() {
 		this.errorSink.asFlux().subscribe(e -> {
 			this.errorHandler.accept(e);
-			// System.err.println(e); TODO: log the error
 		});
 	}
 
@@ -93,7 +87,7 @@ protected Sinks.Many<String> getErrorSink() {
 		return errorSink;
 	}
 
-	public void setInboudMessageHandler(Consumer<JSONRPCMessage> inboundMessageHandler) {
+	public void setInboundMessageHandler(Consumer<JSONRPCMessage> inboundMessageHandler) {
 		this.inboundMessageHandler = inboundMessageHandler;
 	}
 

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/DefaultMcpSession.java

@@ -54,7 +54,7 @@ public DefaultMcpSession(Duration requestTimeout, ObjectMapper objectMapper, Mcp
 		this.objectMapper = objectMapper;
 		this.transport = transport;
 
-		this.transport.setInboudMessageHandler(message -> {
+		this.transport.setInboundMessageHandler(message -> {
 
 			if (message instanceof McpSchema.JSONRPCResponse response) {
 				var sink = pendingResponses.remove(response.id());
@@ -80,33 +80,27 @@ else if (message instanceof McpSchema.JSONRPCNotification notification) {
 
 	@Override
 	public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReference<T> typeRef) {
+		// TODO: UUID API is blocking. Consider non-blocking alternatives to generate
+		// the ID.
 		String requestId = UUID.randomUUID().toString();
 
 		return Mono.<McpSchema.JSONRPCResponse>create(sink -> {
 			this.pendingResponses.put(requestId, sink);
 			McpSchema.JSONRPCRequest jsonrpcRequest = new McpSchema.JSONRPCRequest(McpSchema.JSONRPC_VERSION, method,
 					requestId, requestParams);
-			try {
-				// TODO: This is non-blocking, but it's actually a synchronous call,
-				// perhaps there's no need to make it return Mono?
-				this.transport.sendMessage(jsonrpcRequest)
-					// TODO: It's most efficient to create a dedicated
-					// Subscriber here
-					.subscribe(v -> {
-					}, e -> {
-						this.pendingResponses.remove(requestId);
-						sink.error(e);
-					});
-			}
-			catch (Exception e) {
-				sink.error(e);
-			}
+			this.transport.sendMessage(jsonrpcRequest)
+				// TODO: It's most efficient to create a dedicated Subscriber here
+				.subscribe(v -> {
+				}, e -> {
+					this.pendingResponses.remove(requestId);
+					sink.error(e);
+				});
 		}).timeout(this.requestTimeout).handle((jsonRpcResponse, s) -> {
 			if (jsonRpcResponse.error() != null) {
 				s.error(new McpError(jsonRpcResponse.error()));
 			}
 			else {
-				if (typeRef.getType().getTypeName().equals("java.lang.Void")) {
+				if (typeRef.getType().equals(Void.class)) {
 					s.complete();
 				}
 				else {
@@ -120,19 +114,17 @@ public <T> Mono<T> sendRequest(String method, Object requestParams, TypeReferenc
 	public Mono<Void> sendNotification(String method, Map<String, Object> params) {
 		McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
 				method, params);
-		try {
-			this.transport.sendMessage(jsonrpcNotification);
-		}
-		catch (Exception e) {
-			return Mono.error(new McpError(e));
-		}
-		return Mono.empty();
+		return this.transport.sendMessage(jsonrpcNotification);
+	}
+
+	@Override
+	public Mono<Void> closeGracefully() {
+		return transport.closeGracefully();
 	}
 
 	@Override
-	public Mono<Void> closeGracefully(Duration timeout) {
-		// TODO handle the timeout in transport
-		return Mono.fromRunnable(this.transport::close);
+	public void close() {
+		transport.close();
 	}
 
 }

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSchema.java

@@ -164,17 +164,24 @@ public record Implementation(// @formatter:off
 	} // @formatter:on
 
 	// Existing Enums and Base Types (from previous implementation)
-	public enum Role {
+	public enum Role {// @formatter:off
 
-		USER, ASSISTANT
-
-	}
+		@JsonProperty("user") USER,
+		@JsonProperty("assistant") ASSISTANT
+	}// @formatter:on
 
-	public enum LoggingLevel {
+	public enum LoggingLevel {// @formatter:off
 
-		DEBUG, INFO, NOTICE, WARNING, ERROR, CRITICAL, ALERT, EMERGENCY
+		@JsonProperty("debug") DEBUG,
+		@JsonProperty("info") INFO,
+		@JsonProperty("notice") NOTICE,
+		@JsonProperty("warning") WARNING,
+		@JsonProperty("error") ERROR,
+		@JsonProperty("critical") CRITICAL,
+		@JsonProperty("alert") ALERT,
+		@JsonProperty("emergency") EMERGENCY
 
-	}
+	} // @formatter:on
 
 	// ---------------------------
 	// Resource Interfaces
@@ -448,20 +455,22 @@ public record CreateMessageRequest(List<SamplingMessage> messages, ModelPreferen
 			String systemPrompt, ContextInclusionStrategy includeContext, Double temperature, int maxTokens,
 			List<String> stopSequences, Map<String, Object> metadata) implements Request {
 
-		public enum ContextInclusionStrategy {
-
-			NONE, THIS_SERVER, ALL_SERVERS
-
-		}
+		public enum ContextInclusionStrategy {// @formatter:off
+			@JsonProperty("none") NONE,
+			@JsonProperty("this_server") THIS_SERVER,
+			@JsonProperty("all_server") ALL_SERVERS
+		}	// @formatter:on
 	}
 
 	@JsonInclude(JsonInclude.Include.NON_ABSENT)
 	public record CreateMessageResult(Role role, Content content, String model, StopReason stopReason) {
 		public enum StopReason {
 
-			END_TURN, STOP_SEQUENCE, MAX_TOKENS
-
-		}
+		// @formatter:off
+			@JsonProperty("end_turn") END_TURN,
+			@JsonProperty("stop_sequence") STOP_SEQUENCE,
+			@JsonProperty("max_tokens") MAX_TOKENS
+		} // @formatter:on
 	}
 
 	// ---------------------------

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpSession.java

@@ -81,15 +81,14 @@ default Mono<Void> sendNotification(String method) {
 	Mono<Void> sendNotification(String method, Map<String, Object> params);
 
 	/**
-	 * Gracefully closes the session with a specified timeout.
-	 *
-	 * <p>
-	 * This method attempts to close the session cleanly, waiting for any pending
-	 * operations to complete up to the specified timeout duration.
-	 * </p>
-	 * @param timeout the maximum duration to wait for the session to close
-	 * @return a Mono that completes when the session has been closed
+	 * Closes the session and releases any associated resources asynchronously.
+	 * @return a {@link Mono<Void>} that completes when the session has been closed.
+	 */
+	Mono<Void> closeGracefully();
+
+	/**
+	 * Closes the session and releases any associated resources.
 	 */
-	Mono<Void> closeGracefully(Duration timeout);
+	void close();
 
 }

spring-ai-mcp-core/src/main/java/org/springframework/ai/mcp/spec/McpTransport.java

@@ -91,7 +91,7 @@ default void close() {
 	 * @param inboundMessageHandler a consumer that processes incoming
 	 * {@link JSONRPCMessage}s
 	 */
-	void setInboudMessageHandler(Consumer<JSONRPCMessage> inboundMessageHandler);
+	void setInboundMessageHandler(Consumer<JSONRPCMessage> inboundMessageHandler);
 
 	/**
 	 * Sets the handler for processing transport-level errors.