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

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

@@ -44,7 +44,7 @@ public class McpSyncClient implements AutoCloseable {
 
 	@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.io.IOException;
 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 class StdioServerTransport extends AbstractMcpTransport {
 
 	@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 class StdioServerTransport extends AbstractMcpTransport {
 
 		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 @@ import org.springframework.ai.mcp.spec.McpSchema.JSONRPCMessage;
  * @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 abstract class AbstractMcpTransport implements McpTransport {
 		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 @@ public abstract class AbstractMcpTransport implements McpTransport {
 	private void handleIncomingErrors() {
 		this.errorSink.asFlux().subscribe(e -> {
 			this.errorHandler.accept(e);
-			// System.err.println(e); TODO: log the error
 		});
 	}
 
@@ -93,7 +87,7 @@ public abstract class AbstractMcpTransport implements McpTransport {
 		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 class DefaultMcpSession implements McpSession {
 		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 @@ public class DefaultMcpSession implements McpSession {
 
 	@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 {
+			this.transport.sendMessage(jsonrpcRequest)
-				// TODO: This is non-blocking, but it's actually a synchronous call,
+				// TODO: It's most efficient to create a dedicated Subscriber here
-				// perhaps there's no need to make it return Mono?
+				.subscribe(v -> {
-				this.transport.sendMessage(jsonrpcRequest)
+				}, e -> {
-					// TODO: It's most efficient to create a dedicated
+					this.pendingResponses.remove(requestId);
-					// Subscriber here
+					sink.error(e);
-					.subscribe(v -> {
+				});
-					}, e -> {
-						this.pendingResponses.remove(requestId);
-						sink.error(e);
-					});
-			}
-			catch (Exception e) {
-				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 class DefaultMcpSession implements McpSession {
 	public Mono<Void> sendNotification(String method, Map<String, Object> params) {
 		McpSchema.JSONRPCNotification jsonrpcNotification = new McpSchema.JSONRPCNotification(McpSchema.JSONRPC_VERSION,
 				method, params);
-		try {
+		return this.transport.sendMessage(jsonrpcNotification);
-			this.transport.sendMessage(jsonrpcNotification);
-		}
-		catch (Exception e) {
-			return Mono.error(new McpError(e));
-		}
-		return Mono.empty();
 	}
 
 	@Override
-	public Mono<Void> closeGracefully(Duration timeout) {
+	public Mono<Void> closeGracefully() {
-		// TODO handle the timeout in transport
+		return transport.closeGracefully();
-		return Mono.fromRunnable(this.transport::close);
+	}
+
+	@Override
+	public void close() {
+		transport.close();
 	}
 
 }

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

@@ -164,17 +164,24 @@ public class McpSchema {
 	} // @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 {// @formatter:off
 
-	public enum LoggingLevel {
+		@JsonProperty("debug") DEBUG,
+		@JsonProperty("info") INFO,
+		@JsonProperty("notice") NOTICE,
+		@JsonProperty("warning") WARNING,
+		@JsonProperty("error") ERROR,
+		@JsonProperty("critical") CRITICAL,
+		@JsonProperty("alert") ALERT,
+		@JsonProperty("emergency") EMERGENCY
 
-		DEBUG, INFO, NOTICE, WARNING, ERROR, CRITICAL, ALERT, EMERGENCY
+	} // @formatter:on
-
-	}
 
 	// ---------------------------
 	// Resource Interfaces
@@ -448,20 +455,22 @@ public class McpSchema {
 			String systemPrompt, ContextInclusionStrategy includeContext, Double temperature, int maxTokens,
 			List<String> stopSequences, Map<String, Object> metadata) implements Request {
 
-		public enum ContextInclusionStrategy {
+		public enum ContextInclusionStrategy {// @formatter:off
-
+			@JsonProperty("none") NONE,
-			NONE, THIS_SERVER, ALL_SERVERS
+			@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 @@ public interface McpSession {
 	Mono<Void> sendNotification(String method, Map<String, Object> params);
 
 	/**
-	 * Gracefully closes the session with a specified timeout.
+	 * Closes the session and releases any associated resources asynchronously.
-	 *
+	 * @return a {@link Mono<Void>} that completes when the session has been closed.
-	 * <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
 	 */
-	Mono<Void> closeGracefully(Duration timeout);
+	Mono<Void> closeGracefully();
+
+	/**
+	 * Closes the session and releases any associated resources.
+	 */
+	void close();
 
 }

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

@@ -91,7 +91,7 @@ public interface McpTransport {
 	 * @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.

spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpAsyncClientTests.java

@@ -24,6 +24,7 @@ import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
+import reactor.test.StepVerifier;
 
 import org.springframework.ai.mcp.client.stdio.ServerParameters;
 import org.springframework.ai.mcp.client.stdio.StdioServerTransport;
@@ -41,8 +42,9 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
  * Unit tests for MCP Client Session functionality.
  *
  * @author Christian Tzolov
- * @since 1.0.0
+ * @author Dariusz Jędrzejczyk
  */
+@Timeout(15) // Giving extra time beyond the client timeout
 class McpAsyncClientTests {
 
 	private McpAsyncClient mcpAsyncClient;
@@ -61,15 +63,16 @@ class McpAsyncClientTests {
 
 		assertThatCode(() -> {
 			mcpAsyncClient = McpClient.async(new StdioServerTransport(stdioParams), TIMEOUT, new ObjectMapper());
-			mcpAsyncClient.initialize();
+			mcpAsyncClient.initialize().block(Duration.ofSeconds(10));
 		}).doesNotThrowAnyException();
 	}
 
 	@AfterEach
 	void tearDown() {
-		// if (mcpSyncClient != null) {
+		if (mcpAsyncClient != null) {
-		// assertThatCode(() -> mcpSyncClient.close()).doesNotThrowAnyException();
+			assertThatCode(() -> mcpAsyncClient.closeGracefully().block(Duration.ofSeconds(10)))
-		// }
+				.doesNotThrowAnyException();
+		}
 	}
 
 	@Test
@@ -85,38 +88,34 @@ class McpAsyncClientTests {
 	}
 
 	@Test
-	@Timeout(15) // Giving extra time beyond the client timeout
 	void testListTools() {
-		mcpAsyncClient.listTools(null).subscribe(result -> {
+		StepVerifier.create(mcpAsyncClient.listTools(null)).consumeNextWith(result -> {
 			assertThat(result.tools()).isNotNull().isNotEmpty();
 
 			Tool firstTool = result.tools().get(0);
 			assertThat(firstTool.name()).isNotNull();
 			assertThat(firstTool.description()).isNotNull();
-		});
+		}).verifyComplete();
 	}
 
 	@Test
-	@Timeout(15)
 	void testPing() {
 		assertThatCode(() -> mcpAsyncClient.ping().block()).doesNotThrowAnyException();
 	}
 
 	@Test
-	@Timeout(15)
 	void testCallTool() {
 		CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE));
 
-		mcpAsyncClient.callTool(callToolRequest).subscribe(callToolResult -> {
+		StepVerifier.create(mcpAsyncClient.callTool(callToolRequest)).consumeNextWith(callToolResult -> {
 			assertThat(callToolResult).isNotNull().satisfies(result -> {
 				assertThat(result.content()).isNotNull();
 				assertThat(result.isError()).isNull();
 			});
-		});
+		}).verifyComplete();
 	}
 
 	@Test
-	@Timeout(15)
 	void testCallToolWithInvalidTool() {
 		CallToolRequest invalidRequest = new CallToolRequest("nonexistent_tool", Map.of("message", TEST_MESSAGE));
 
@@ -124,15 +123,13 @@ class McpAsyncClientTests {
 	}
 
 	@Test
-	@Timeout(15)
 	void testRootsListChanged() {
 		assertThatCode(() -> mcpAsyncClient.sendRootsListChanged().block()).doesNotThrowAnyException();
 	}
 
 	@Test
-	@Timeout(15)
 	void testListResources() {
-		mcpAsyncClient.listResources(null).subscribe(resources -> {
+		StepVerifier.create(mcpAsyncClient.listResources(null)).consumeNextWith(resources -> {
 			assertThat(resources).isNotNull().satisfies(result -> {
 				assertThat(result.resources()).isNotNull();
 
@@ -142,7 +139,7 @@ class McpAsyncClientTests {
 					assertThat(firstResource.name()).isNotNull();
 				}
 			});
-		});
+		}).verifyComplete();
 	}
 
 	@Test
@@ -151,9 +148,8 @@ class McpAsyncClientTests {
 	}
 
 	@Test
-	@Timeout(15)
 	void testListPrompts() {
-		mcpAsyncClient.listPrompts(null).subscribe(prompts -> {
+		StepVerifier.create(mcpAsyncClient.listPrompts(null)).consumeNextWith(prompts -> {
 			assertThat(prompts).isNotNull().satisfies(result -> {
 				assertThat(result.prompts()).isNotNull();
 
@@ -163,18 +159,19 @@ class McpAsyncClientTests {
 					assertThat(firstPrompt.description()).isNotNull();
 				}
 			});
-		});
+		}).verifyComplete();
 	}
 
 	@Test
-	@Timeout(15)
 	void testGetPrompt() {
-		mcpAsyncClient.getPrompt(new GetPromptRequest("simple_prompt", Map.of())).subscribe(prompt -> {
+		StepVerifier.create(mcpAsyncClient.getPrompt(new GetPromptRequest("simple_prompt", Map.of())))
-			assertThat(prompt).isNotNull().satisfies(result -> {
+			.consumeNextWith(prompt -> {
-				assertThat(result.description()).isNotNull();
+				assertThat(prompt).isNotNull().satisfies(result -> {
-				assertThat(result.messages()).isNotEmpty();
+					assertThat(result.messages()).isNotEmpty();
-			});
+					assertThat(result.messages()).hasSize(1);
-		});
+				});
+			})
+			.verifyComplete();
 	}
 
 }

spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/client/McpSyncClientTests.java

@@ -43,8 +43,9 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
  * Unit tests for MCP Client Session functionality.
  *
  * @author Christian Tzolov
- * @since 1.0.0
+ * @author Dariusz Jędrzejczyk
  */
+@Timeout(15) // Giving extra time beyond the client timeout
 class McpSyncClientTests {
 
 	private McpSyncClient mcpSyncClient;
@@ -87,7 +88,6 @@ class McpSyncClientTests {
 	}
 
 	@Test
-	@Timeout(15) // Giving extra time beyond the client timeout
 	void testListTools() {
 		ListToolsResult tools = mcpSyncClient.listTools(null);
 
@@ -101,7 +101,6 @@ class McpSyncClientTests {
 	}
 
 	@Test
-	@Timeout(15) // Giving extra time beyond the client timeout
 	void testCallTools() {
 		CallToolResult toolResult = mcpSyncClient.callTool(new CallToolRequest("add", Map.of("a", 3, "b", 4)));
 
@@ -120,13 +119,11 @@ class McpSyncClientTests {
 	}
 
 	@Test
-	@Timeout(15)
 	void testPing() {
 		assertThatCode(() -> mcpSyncClient.ping()).doesNotThrowAnyException();
 	}
 
 	@Test
-	@Timeout(15)
 	void testCallTool() {
 		CallToolRequest callToolRequest = new CallToolRequest("echo", Map.of("message", TEST_MESSAGE));
 
@@ -139,7 +136,6 @@ class McpSyncClientTests {
 	}
 
 	@Test
-	@Timeout(15)
 	void testCallToolWithInvalidTool() {
 		CallToolRequest invalidRequest = new CallToolRequest("nonexistent_tool", Map.of("message", TEST_MESSAGE));
 
@@ -147,13 +143,11 @@ class McpSyncClientTests {
 	}
 
 	@Test
-	@Timeout(15)
 	void testRootsListChanged() {
 		assertThatCode(() -> mcpSyncClient.sendRootsListChanged()).doesNotThrowAnyException();
 	}
 
 	@Test
-	@Timeout(15)
 	void testListResources() {
 		ListResourcesResult resources = mcpSyncClient.listResources(null);
 

spring-ai-mcp-core/src/test/java/org/springframework/ai/mcp/spec/DefaultMcpTransportTests.java

@@ -42,7 +42,7 @@ class DefaultMcpTransportTests {
 		transport = new AbstractMcpTransport(objectMapper) {
 			@Override
 			public Mono<Void> closeGracefully() {
-				return null;
+				return Mono.empty();
 			}
 		};
 	}
@@ -93,7 +93,7 @@ class DefaultMcpTransportTests {
 	@Test
 	void customInboundMessageHandlerShouldReceiveMessages() {
 		AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
-		transport.setInboudMessageHandler(receivedMessage::set);
+		transport.setInboundMessageHandler(receivedMessage::set);
 		transport.start();
 
 		JSONRPCRequest message = new JSONRPCRequest("2.0", "test", 1, null);
@@ -123,7 +123,7 @@ class DefaultMcpTransportTests {
 		AtomicReference<JSONRPCMessage> receivedMessage = new AtomicReference<>();
 		AtomicReference<String> receivedError = new AtomicReference<>();
 
-		transport.setInboudMessageHandler(receivedMessage::set);
+		transport.setInboundMessageHandler(receivedMessage::set);
 		transport.setInboundErrorHandler(receivedError::set);
 		transport.start();