diff --git a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
index 3436fb4f5..f2fcd9738 100644
--- a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
+++ b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/PartitionAwareFunctionWrapper.java
@@ -41,13 +41,14 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 
 	protected final Log logger = LogFactory.getLog(PartitionAwareFunctionWrapper.class);
 
-	private final org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper function;
+	@SuppressWarnings("rawtypes")
+	private final Function function;
 
 	@SuppressWarnings("rawtypes")
 	private final Function<Message, Message> outputMessageEnricher;
 
-	@SuppressWarnings("unchecked")
-	PartitionAwareFunctionWrapper(FunctionInvocationWrapper function, ConfigurableApplicationContext context, ProducerProperties producerProperties) {
+	@SuppressWarnings({ "unchecked", "rawtypes" })
+	PartitionAwareFunctionWrapper(Function function, ConfigurableApplicationContext context, ProducerProperties producerProperties) {
 		this.function = function;
 		if (producerProperties != null && producerProperties.isPartitioned()) {
 			StandardEvaluationContext evaluationContext = ExpressionUtils.createStandardEvaluationContext(context.getBeanFactory());
@@ -65,26 +66,36 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 		}
 	}
 
+	@SuppressWarnings("unchecked")
 	@Override
 	public Object apply(Object input) {
-		if (this.outputMessageEnricher == null) { // to avoid breaking change
-			return this.function.apply(input);
+		if (this.function instanceof FunctionInvocationWrapper && this.outputMessageEnricher != null) {
+			try {
+				return ((FunctionInvocationWrapper) this.function).apply(input, this.outputMessageEnricher);
+			}
+			catch (NoSuchMethodError e) {
+				logger.warn("Versions of spring-cloud-function older then 3.0.2.RELEASE do not support generation of partition information. "
+						+ "Output message will not contain any partition header unless spring-cloud-function dependency is 3.0.2.RELEASE or higher.");
+				return this.function.apply(input);
+			}
 		}
-		try {
-			return this.function.apply(input, this.outputMessageEnricher);
-		}
-		catch (NoSuchMethodError e) {
-			logger.warn("Versions of spring-cloud-function older then 3.0.2.RELEASE do not support generation of partition information. "
-					+ "Output message will not contain any partition header unless spring-cloud-function dependency is 3.0.2.RELEASE or higher.");
+		else {
 			return this.function.apply(input);
 		}
 	}
 
 	@Override
 	public Object get() {
-		if (this.outputMessageEnricher == null) { // to avoid breaking change
-			return this.function.get();
+		if (this.function instanceof FunctionInvocationWrapper) {
+			if (this.outputMessageEnricher != null) {
+				return ((FunctionInvocationWrapper) this.function).get(this.outputMessageEnricher);
+			}
+			else {
+				return ((FunctionInvocationWrapper) this.function).get();
+			}
+		}
+		else {
+			throw new IllegalStateException("THis function is not a Supplier. Call to get() is not allowed");
 		}
-		return this.function.get(this.outputMessageEnricher);
 	}
 }
diff --git a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java
index 90004cb08..66dcb8559 100644
--- a/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java
+++ b/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/StreamBridge.java
@@ -31,7 +31,6 @@ import org.springframework.cloud.function.context.FunctionCatalog;
 import org.springframework.cloud.function.context.FunctionRegistration;
 import org.springframework.cloud.function.context.FunctionRegistry;
 import org.springframework.cloud.function.context.FunctionType;
-import org.springframework.cloud.function.context.catalog.SimpleFunctionRegistry.FunctionInvocationWrapper;
 import org.springframework.cloud.stream.binder.ProducerProperties;
 import org.springframework.cloud.stream.binding.BinderAwareChannelResolver.NewDestinationBindingCallback;
 import org.springframework.cloud.stream.binding.BindingService;
@@ -161,7 +160,7 @@ public final class StreamBridge implements SmartInitializingSingleton {
 						: this.functionCatalog.lookup(STREAM_BRIDGE_FUNC_NAME, outputContentType.toString());
 
 		if (producerProperties != null && producerProperties.isPartitioned()) {
-			functionToInvoke = new PartitionAwareFunctionWrapper((FunctionInvocationWrapper) functionToInvoke, this.applicationContext, producerProperties);
+			functionToInvoke = new PartitionAwareFunctionWrapper(functionToInvoke, this.applicationContext, producerProperties);
 		}
 		// this function is a pass through and is only required to force output conversion if necessary on SCF side.
 		Message<byte[]> resultMessage = (Message<byte[]>) functionToInvoke.apply(data);