GH-1294 Additional changes related to removal of TARGET_PROTOCOL

Author: olegz

spring-cloud-function-context/src/main/java/org/springframework/cloud/function/cloudevent/CloudEventsFunctionInvocationHelper.java

@@ -71,7 +71,7 @@ public class CloudEventsFunctionInvocationHelper implements FunctionInvocationHe
 
 	@Override
 	public boolean isRetainOutputAsMessage(Message<?> message) {
-		return (message.getHeaders().containsKey(MessageUtils.MESSAGE_TYPE)
+		return message.getHeaders().containsKey("partitionKey") || (message.getHeaders().containsKey(MessageUtils.MESSAGE_TYPE)
 			&& message.getHeaders().get(MessageUtils.MESSAGE_TYPE).equals(CloudEventMessageUtils.CLOUDEVENT_VALUE));
 	}
 

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

@@ -1235,7 +1235,9 @@ else if (isExtractPayload((Message<?>) convertedOutput, type)) {
 				}
 				else {
 					enrichedMessage = MessageBuilder.withPayload(convertedOutput)
-						.setHeader(MessageHeaders.CONTENT_TYPE, contentType[0]).build();
+						.setHeader(MessageHeaders.CONTENT_TYPE, contentType[0])
+						.copyHeaders(((Message) output).getHeaders())
+						.build();
 				}
 				return messageConverter.toMessage(enrichedMessage.getPayload(), enrichedMessage.getHeaders());
 			}