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 28c8efce6..6cf0fb41f 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,16 +41,14 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 
 	protected final Log logger = LogFactory.getLog(PartitionAwareFunctionWrapper.class);
 
-	private final FunctionInvocationWrapper function;
-
-	private Function<Object, Message> enhancerFunction;
+	@SuppressWarnings("rawtypes")
+	private final Function function;
 
 	@SuppressWarnings("rawtypes")
 	private final Function<Object, Message> outputMessageEnricher;
 
-	PartitionAwareFunctionWrapper(FunctionInvocationWrapper function, ConfigurableApplicationContext context, ProducerProperties producerProperties) {
+	PartitionAwareFunctionWrapper(Function<?, ?> function, ConfigurableApplicationContext context, ProducerProperties producerProperties) {
 		this.function = function;
-		this.enhancerFunction = function.getEnhancer();
 
 		if (producerProperties != null && producerProperties.isPartitioned()) {
 			StandardEvaluationContext evaluationContext = ExpressionUtils.createStandardEvaluationContext(context.getBeanFactory());
@@ -71,6 +69,7 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 		}
 	}
 
+	@SuppressWarnings("unchecked")
 	@Override
 	public Object apply(Object input) {
 		this.setEnhancerIfNecessary();
@@ -79,11 +78,16 @@ class PartitionAwareFunctionWrapper implements Function<Object, Object>, Supplie
 
 	@Override
 	public Object get() {
-		this.setEnhancerIfNecessary();
-		return this.function.get();
+		if (this.function instanceof FunctionInvocationWrapper) {
+			this.setEnhancerIfNecessary();
+			return ((FunctionInvocationWrapper) this.function).get();
+		}
+		throw new IllegalStateException("Call to get() is not allowed since this function is not a Supplier.");
 	}
 
 	private void setEnhancerIfNecessary() {
-		this.function.setEnhancer(this.outputMessageEnricher);
+		if (this.function instanceof FunctionInvocationWrapper) {
+			((FunctionInvocationWrapper) this.function).setEnhancer(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 882ce6cf7..0e4ef8e13 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
@@ -29,7 +29,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.Binder;
 import org.springframework.cloud.stream.binder.BinderFactory;
 import org.springframework.cloud.stream.binder.ProducerProperties;
@@ -203,11 +202,11 @@ public final class StreamBridge implements SmartInitializingSingleton {
 		boolean skipConversion = producerProperties.isUseNativeEncoding();
 
 		Function<Object, Object> functionToInvoke = skipConversion
-				? v -> v instanceof Message ? v :  MessageBuilder.withPayload(v).build()
+				? (v -> v instanceof Message ? v :  MessageBuilder.withPayload(v).build())
 						: 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);