Fix POJO functions to return Message consistently

POJO functions now return Message when input is Message, maintaining
consistency with regular Function implementations.

Fixes gh-1307

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

@@ -287,6 +287,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;
 	}
 
@@ -439,6 +450,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;
@@ -489,6 +502,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;
 		}
@@ -1245,6 +1266,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,34 @@ 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");
+	}
+
 
 	@EnableAutoConfiguration
 	@Configuration(proxyBeanMethods = false)