feat: add integration tests and post-processor for handling Flux return types in MCP tools

Signed-off-by: liugddx <[email protected]>

Author: liugddx

auto-configurations/mcp/spring-ai-autoconfigure-mcp-server-common/src/main/java/org/springframework/ai/mcp/server/common/autoconfigure/annotations/FluxToolSpecificationPostProcessor.java

@@ -16,116 +16,230 @@
 
 package org.springframework.ai.mcp.server.common.autoconfigure.annotations;
 
+import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.function.BiFunction;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import io.modelcontextprotocol.server.McpStatelessServerFeatures;
 import io.modelcontextprotocol.spec.McpSchema;
 import io.modelcontextprotocol.spec.McpTransportContext;
-import org.reactivestreams.Publisher;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springaicommunity.mcp.annotation.McpTool;
+import org.springaicommunity.mcp.annotation.McpToolParam;
 import reactor.core.publisher.Flux;
 import reactor.core.publisher.Mono;
 
+import org.springframework.util.ReflectionUtils;
+
 /**
- * Post-processor that attempts to wrap AsyncToolSpecifications to handle Flux return
- * types properly by collecting all elements.
- * 
+ * Post-processor that wraps AsyncToolSpecifications to handle Flux return types properly
+ * by collecting all elements before serialization.
+ *
  * <p>
- * <strong>NOTE:</strong> This class demonstrates the intended fix for Issue #4542 where
- * Flux-returning tool methods only return the first element. However, the actual fix
- * cannot be applied at this level because the Flux has already been improperly consumed
- * by the time the AsyncToolSpecification's callHandler is invoked.
- * 
+ * <strong>Background:</strong> This class fixes Issue #4542 where Flux-returning @McpTool
+ * methods only return the first element. The root cause is in the external {@code
+ * org.springaicommunity.mcp.provider.tool.AsyncStatelessMcpToolProvider} library, which
+ * treats Flux as a single-value Publisher and only takes the first element.
+ *
  * <p>
- * The real fix needs to be in the {@code org.springaicommunity.mcp.provider.tool.AsyncStatelessMcpToolProvider}
- * class (from the external {@code mcp-annotations} library), where the annotated method
- * is invoked and its return value is processed. That code needs to check if the return
- * value is a {@code Flux<T>} and call {@code .collectList()} before converting it to a
- * {@code CallToolResult}.
- * 
+ * <strong>Solution:</strong> This post-processor intercepts tool specifications and wraps
+ * their call handlers. When a method returns a Flux, it collects all elements into a list
+ * before passing the result to the MCP serialization layer.
+ *
  * <p>
- * <strong>Workaround:</strong> Users should return {@code Mono<List<T>>} instead of
- * {@code Flux<T>} from their {@code @McpTool} methods. See
- * {@code FLUX-RETURN-TYPE-WORKAROUND.md} for details.
- * 
+ * <strong>Note:</strong> Users can also work around this issue by returning {@code
+ * Mono<List<T>>} instead of {@code Flux<T>} from their {@code @McpTool} methods.
+ *
+ * @author liugddx
  * @since 1.1.0
  * @see <a href="https://github.com/spring-projects/spring-ai/issues/4542">Issue #4542</a>
  */
-public class FluxToolSpecificationPostProcessor {
+public final class FluxToolSpecificationPostProcessor {
+
+	private static final Logger logger = LoggerFactory.getLogger(FluxToolSpecificationPostProcessor.class);
+
+	private static final ObjectMapper objectMapper = new ObjectMapper();
+
+	private FluxToolSpecificationPostProcessor() {
+		// Utility class - no instances allowed
+	}
 
 	/**
 	 * Wraps tool specifications to properly handle Flux return types by collecting all
 	 * elements into a list.
 	 * @param originalSpecs the original tool specifications from the annotation provider
+	 * @param toolBeans the bean objects containing @McpTool annotated methods
 	 * @return wrapped tool specifications that properly collect Flux elements
 	 */
 	public static List<McpStatelessServerFeatures.AsyncToolSpecification> processToolSpecifications(
-			List<McpStatelessServerFeatures.AsyncToolSpecification> originalSpecs) {
+			List<McpStatelessServerFeatures.AsyncToolSpecification> originalSpecs, List<Object> toolBeans) {
 
-		List<McpStatelessServerFeatures.AsyncToolSpecification> wrappedSpecs = new ArrayList<>();
+		List<McpStatelessServerFeatures.AsyncToolSpecification> processedSpecs = new ArrayList<>();
 
 		for (McpStatelessServerFeatures.AsyncToolSpecification spec : originalSpecs) {
-			McpStatelessServerFeatures.AsyncToolSpecification wrappedSpec = wrapToolSpecification(spec);
-			wrappedSpecs.add(wrappedSpec);
+			ToolMethodInfo methodInfo = findToolMethod(toolBeans, spec.tool().name());
+			if (methodInfo != null && methodInfo.returnsFlux()) {
+				logger.info("Detected Flux return type for MCP tool '{}', applying collection wrapper",
+						spec.tool().name());
+				McpStatelessServerFeatures.AsyncToolSpecification wrappedSpec = wrapToolSpecificationForFlux(spec,
+						methodInfo);
+				processedSpecs.add(wrappedSpec);
+			}
+			else {
+				processedSpecs.add(spec);
+			}
 		}
 
-		return wrappedSpecs;
+		return processedSpecs;
 	}
 
-	private static McpStatelessServerFeatures.AsyncToolSpecification wrapToolSpecification(
-			McpStatelessServerFeatures.AsyncToolSpecification original) {
+	/**
+	 * Finds the method annotated with @McpTool that matches the given tool name.
+	 * @param toolBeans the bean objects containing @McpTool annotated methods
+	 * @param toolName the name of the tool to find
+	 * @return the ToolMethodInfo object, or null if not found
+	 */
+	private static ToolMethodInfo findToolMethod(List<Object> toolBeans, String toolName) {
+		for (Object bean : toolBeans) {
+			Class<?> clazz = bean.getClass();
+			Method[] methods = ReflectionUtils.getAllDeclaredMethods(clazz);
+			for (Method method : methods) {
+				McpTool annotation = method.getAnnotation(McpTool.class);
+				if (annotation != null && annotation.name().equals(toolName)) {
+					return new ToolMethodInfo(bean, method);
+				}
+			}
+		}
+		return null;
+	}
+
+	/**
+	 * Wraps a tool specification to collect all Flux elements before serialization.
+	 * @param original the original tool specification
+	 * @param methodInfo the method information including bean and method
+	 * @return the wrapped tool specification
+	 */
+	private static McpStatelessServerFeatures.AsyncToolSpecification wrapToolSpecificationForFlux(
+			McpStatelessServerFeatures.AsyncToolSpecification original, ToolMethodInfo methodInfo) {
 
 		BiFunction<McpTransportContext, McpSchema.CallToolRequest, Mono<McpSchema.CallToolResult>> originalHandler = original
 			.callHandler();
 
 		BiFunction<McpTransportContext, McpSchema.CallToolRequest, Mono<McpSchema.CallToolResult>> wrappedHandler = (
 				context, request) -> {
-			Mono<McpSchema.CallToolResult> result = originalHandler.apply(context, request);
-
-			// The issue is that if the underlying implementation returns a Flux but
-			// only emits the first element, we need to ensure all elements are
-			// collected.
-			// However, at this level, we receive a Mono<CallToolResult>, so the damage
-			// is already done.
-			// The fix needs to be in the org.springaicommunity.mcp library itself.
-
-			return result;
+			try {
+				// Invoke the method directly to get access to the Flux
+				Object[] args = buildMethodArguments(methodInfo.method(), request.arguments());
+				Object result = ReflectionUtils.invokeMethod(methodInfo.method(), methodInfo.bean(), args);
+
+				if (result instanceof Flux) {
+					// Collect all Flux elements into a list
+					Flux<?> flux = (Flux<?>) result;
+					return flux.collectList().flatMap(list -> {
+						// Serialize the list to JSON
+						try {
+							String jsonContent = objectMapper.writeValueAsString(list);
+							return Mono.just(new McpSchema.CallToolResult(
+									List.of(new McpSchema.TextContent(jsonContent)), false));
+						}
+						catch (Exception e) {
+							logger.error("Failed to serialize Flux result for tool '{}'", original.tool().name(), e);
+							return Mono.just(new McpSchema.CallToolResult(
+									List.of(new McpSchema.TextContent("Error: " + e.getMessage())), true));
+						}
+					});
+				}
+				else {
+					// Fall back to original handler for non-Flux results
+					return originalHandler.apply(context, request);
+				}
+			}
+			catch (Exception e) {
+				logger.error("Failed to invoke tool method '{}'", original.tool().name(), e);
+				return Mono.just(
+						new McpSchema.CallToolResult(List.of(new McpSchema.TextContent("Error: " + e.getMessage())),
+								true));
+			}
 		};
 
 		return new McpStatelessServerFeatures.AsyncToolSpecification(original.tool(), wrappedHandler);
 	}
 
 	/**
-	 * Helper method to detect if a Publisher might be a Flux that needs special
-	 * handling.
-	 * @param publisher the publisher to check
-	 * @return true if the publisher is a Flux (multi-element stream)
+	 * Builds method arguments from the request arguments map.
+	 * @param method the method to invoke
+	 * @param requestArgs the arguments from the CallToolRequest
+	 * @return array of method arguments
 	 */
-	private static boolean isFlux(Publisher<?> publisher) {
-		return publisher instanceof Flux;
+	private static Object[] buildMethodArguments(Method method, Map<String, Object> requestArgs) {
+		java.lang.reflect.Parameter[] parameters = method.getParameters();
+		Object[] args = new Object[parameters.length];
+
+		for (int i = 0; i < parameters.length; i++) {
+			java.lang.reflect.Parameter param = parameters[i];
+			McpToolParam paramAnnotation = param.getAnnotation(McpToolParam.class);
+
+			if (paramAnnotation != null) {
+				String paramName = paramAnnotation.name().isEmpty() ? param.getName() : paramAnnotation.name();
+				Object value = requestArgs.get(paramName);
+
+				// Type conversion if needed
+				if (value != null) {
+					args[i] = objectMapper.convertValue(value, param.getType());
+				}
+				else if (!paramAnnotation.required()) {
+					args[i] = null;
+				}
+				else {
+					throw new IllegalArgumentException("Required parameter '" + paramName + "' is missing");
+				}
+			}
+			else {
+				// Try to match by parameter name
+				Object value = requestArgs.get(param.getName());
+				if (value != null) {
+					args[i] = objectMapper.convertValue(value, param.getType());
+				}
+				else {
+					args[i] = null;
+				}
+			}
+		}
+
+		return args;
 	}
 
 	/**
-	 * Attempts to collect all elements from a Flux into a list. If the publisher is
-	 * already a Mono, returns it as-is.
-	 * @param publisher the publisher
-	 * @return a Mono containing either the single value or a list of all values
+	 * Holds information about a tool method.
 	 */
-	@SuppressWarnings("unchecked")
-	private static <T> Mono<T> collectIfFlux(Publisher<T> publisher) {
-		if (publisher instanceof Flux) {
-			// Collect all elements into a list
-			return ((Flux<T>) publisher).collectList().map(list -> (T) list);
+	private static class ToolMethodInfo {
+
+		private final Object bean;
+
+		private final Method method;
+
+		ToolMethodInfo(Object bean, Method method) {
+			this.bean = bean;
+			this.method = method;
+			ReflectionUtils.makeAccessible(method);
 		}
-		else if (publisher instanceof Mono) {
-			return (Mono<T>) publisher;
+
+		Object bean() {
+			return this.bean;
 		}
-		else {
-			// Generic Publisher - convert to Flux and collect
-			return Flux.from(publisher).collectList().map(list -> (T) list);
+
+		Method method() {
+			return this.method;
 		}
+
+		boolean returnsFlux() {
+			return Flux.class.isAssignableFrom(this.method.getReturnType());
+		}
+
 	}
 
 }
-

auto-configurations/mcp/spring-ai-autoconfigure-mcp-server-common/src/main/java/org/springframework/ai/mcp/server/common/autoconfigure/annotations/StatelessServerSpecificationFactoryAutoConfiguration.java

@@ -127,8 +127,12 @@ public List<McpStatelessServerFeatures.AsyncCompletionSpecification> completionS
 		@Bean
 		public List<McpStatelessServerFeatures.AsyncToolSpecification> toolSpecs(
 				ServerMcpAnnotatedBeans beansWithMcpMethodAnnotations) {
-			return AsyncMcpAnnotationProviders
-				.statelessToolSpecifications(beansWithMcpMethodAnnotations.getBeansByAnnotation(McpTool.class));
+			List<Object> toolBeans = beansWithMcpMethodAnnotations.getBeansByAnnotation(McpTool.class);
+			List<McpStatelessServerFeatures.AsyncToolSpecification> originalSpecs = AsyncMcpAnnotationProviders
+				.statelessToolSpecifications(toolBeans);
+
+			// Apply post-processing to handle Flux return types (Issue #4542)
+			return FluxToolSpecificationPostProcessor.processToolSpecifications(originalSpecs, toolBeans);
 		}
 
 	}