Fix POJO functions to return Message consistently

POJO functions now return Message when input is Message, maintaining
consistency with regular Function implementations. This ensures headers
are preserved in POJO functions just like regular functions.

Changes:
- Add isPojoFunction flag to identify POJO functions
- Mark POJO functions during discovery
- Re-wrap output in Message for POJO functions when input is Message
- Add test verifying Message return with header preservation
- Add test verifying plain String input returns plain String output

Fixes gh-1307

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

Resolves #1311

Author: JunggiKim

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwareFunctionRegistry.java

@@ -21,6 +21,7 @@
 import java.lang.reflect.Type;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Set;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
@@ -171,6 +172,10 @@ else if (this.isFunctionPojo(functionCandidate, functionName)) {
 								Method functionalMethod = FunctionTypeUtils.discoverFunctionalMethod(functionCandidate.getClass());
 								functionCandidate = this.proxyTarget(functionCandidate, functionalMethod);
 								functionType = FunctionTypeUtils.fromFunctionMethod(functionalMethod);
+								// GH-1307: Mark this as a POJO function for special handling
+								functionRegistration = new FunctionRegistration(functionCandidate, functionName)
+										.type(functionType)
+										.properties(Collections.singletonMap("isPojoFunction", "true"));
 							}
 							else if (this.isSpecialFunctionRegistration(functionNames, functionName)) {
 								functionRegistration = this.applicationContext

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/context/catalog/SimpleFunctionRegistry.java

@@ -291,6 +291,17 @@ private FunctionInvocationWrapper findFunctionInFunctionRegistrations(String fun
 				// ignore
 			}
 		}
+		// GH-1307: Mark POJO functions for special Message wrapping behavior
+		if (functionRegistration != null &&
+				functionRegistration.getProperties().containsKey("isPojoFunction")) {
+			try {
+				String isPojoValue = functionRegistration.getProperties().get("isPojoFunction");
+				function.setPojoFunction(Boolean.parseBoolean(isPojoValue));
+			}
+			catch (Exception e) {
+				// ignore
+			}
+		}
 		return function;
 	}
 
@@ -443,6 +454,8 @@ public class FunctionInvocationWrapper implements Function<Object, Object>, Cons
 
 		private boolean wrappedBiConsumer;
 
+		private boolean isPojoFunction;
+
 		FunctionInvocationWrapper(String functionDefinition,  Object target, Type inputType, Type outputType) {
 			if (target instanceof PostProcessingFunction) {
 				this.postProcessor = (PostProcessingFunction) target;
@@ -498,6 +511,14 @@ public void setWrappedBiConsumer(boolean wrappedBiConsumer) {
 			this.wrappedBiConsumer = wrappedBiConsumer;
 		}
 
+		public void setPojoFunction(boolean isPojoFunction) {
+			this.isPojoFunction = isPojoFunction;
+		}
+
+		public boolean isPojoFunction() {
+			return this.isPojoFunction;
+		}
+
 		public boolean isSkipOutputConversion() {
 			return skipOutputConversion;
 		}
@@ -1259,6 +1280,14 @@ else if (isExtractPayload((Message<?>) convertedOutput, type)) {
 			}
 
 			if (ObjectUtils.isEmpty(contentType)) {
+				// GH-1307: For POJO functions, wrap output in Message to maintain
+				// consistency with regular functions
+				if (this.isPojoFunction && output instanceof Message
+						&& !(convertedOutput instanceof Message)) {
+					convertedOutput = MessageBuilder.withPayload(convertedOutput)
+						.copyHeaders(((Message) output).getHeaders())
+						.build();
+				}
 				return convertedOutput;
 			}
 

spring-cloud-function-context/src/test/java/org/springframework/cloud/function/context/catalog/BeanFactoryAwarePojoFunctionRegistryTests.java

@@ -83,8 +83,10 @@ public void testWithPojoFunction() {
 		Function<Integer, String> f2conversion = catalog.lookup("myFunctionLike");
 		assertThat(f2conversion.apply(123)).isEqualTo("123");
 
-		Function<Message<String>, String> f2message = catalog.lookup("myFunctionLike");
-		assertThat(f2message.apply(MessageBuilder.withPayload("message").build())).isEqualTo("MESSAGE");
+		// GH-1307: POJO functions now return Message for consistency
+		Function<Message<String>, Message<?>> f2message = catalog.lookup("myFunctionLike");
+		Message<?> messageResult = f2message.apply(MessageBuilder.withPayload("message").build());
+		assertThat(messageResult.getPayload()).isEqualTo("MESSAGE");
 
 		Function<Flux<String>, Flux<String>> f3 = catalog.lookup("myFunctionLike");
 		assertThat(f3.apply(Flux.just("foo")).blockFirst()).isEqualTo("FOO");
@@ -100,6 +102,52 @@ public void testWithPojoFunctionComposition() {
 		assertThat(f1.apply("foo")).isEqualTo("FOO");
 	}
 
+	/**
+	 * GH-1307: POJO function should return Message consistently with regular functions
+	 * when no contentType is specified.
+	 */
+	@Test
+	public void testPojoFunctionReturnsMessageWithoutContentType() {
+		FunctionCatalog catalog = this.configureCatalog();
+
+		// Test POJO function without contentType
+		Function<Message<String>, Object> pojoFunction = catalog.lookup("myFunctionLike");
+		Message<String> input = MessageBuilder.withPayload("test")
+			.setHeader("correlationId", "123")
+			.build();
+
+		Object result = pojoFunction.apply(input);
+
+		// GH-1307: Verify POJO functions return Message for consistency
+		assertThat(result)
+			.as("POJO function should return Message, not plain value when input is Message")
+			.isInstanceOf(Message.class);
+
+		Message<?> messageResult = (Message<?>) result;
+		assertThat(messageResult.getPayload()).isEqualTo("TEST");
+		assertThat(messageResult.getHeaders().get("correlationId"))
+			.as("Headers should be preserved")
+			.isEqualTo("123");
+	}
+
+	/**
+	 * GH-1307: POJO function should NOT wrap output when input is plain String
+	 */
+	@Test
+	public void testPojoFunctionDoesNotWrapPlainStringInput() {
+		FunctionCatalog catalog = this.configureCatalog();
+
+		// GH-1307: POJO function with plain String input should return plain String
+		Function<String, Object> pojoFunction = catalog.lookup("myFunctionLike");
+		Object result = pojoFunction.apply("plainInput");
+
+		// Should return String, not Message
+		assertThat(result)
+			.as("POJO function should return plain String when input is plain String, not wrap in Message")
+			.isInstanceOf(String.class)
+			.isEqualTo("PLAININPUT");
+	}
+
 
 	@EnableAutoConfiguration
 	@Configuration(proxyBeanMethods = false)