Issue-3938 Upgrade the Knative custom functions to have visibility over the returned status code and headers

kogito-workitems/kogito-rest-workitem/src/main/java/org/kogito/workitem/rest/decorators/AbstractParamsDecorator.java

@@ -18,9 +18,11 @@
  */
 package org.kogito.workitem.rest.decorators;
 
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Map;
 import java.util.Map.Entry;
+import java.util.Set;
 
 import org.kie.kogito.internal.process.workitem.KogitoWorkItem;
 
@@ -30,26 +32,36 @@ public abstract class AbstractParamsDecorator implements ParamsDecorator {
 
     @Override
     public void decorate(KogitoWorkItem item, Map<String, Object> parameters, HttpRequest<?> request) {
+        extractHeadersQueries(item, parameters, request);
+    }
+
+    protected String toHeaderKey(String key) {
+        return key;
+    }
+
+    protected String toQueryKey(String key) {
+        return key;
+    }
+
+    protected Set<String> extractHeadersQueries(KogitoWorkItem item, Map<String, Object> parameters, HttpRequest<?> request) {
+        Set<String> consideredParams = new HashSet<>();
+
         Iterator<Entry<String, Object>> iter = parameters.entrySet().iterator();
         while (iter.hasNext()) {
             Entry<String, Object> entry = iter.next();
             String key = entry.getKey();
             if (isHeaderParameter(key)) {
                 request.putHeader(toHeaderKey(key), entry.getValue().toString());
+                consideredParams.add(key);
                 iter.remove();
             } else if (isQueryParameter(key)) {
                 request.addQueryParam(toQueryKey(key), entry.getValue().toString());
+                consideredParams.add(key);
                 iter.remove();
             }
         }
-    }
 
-    protected String toHeaderKey(String key) {
-        return key;
-    }
-
-    protected String toQueryKey(String key) {
-        return key;
+        return consideredParams;
     }
 
     protected abstract boolean isHeaderParameter(String key);

quarkus/addons/knative/serving/deployment/src/main/java/org/kie/kogito/addons/quarkus/knative/serving/deployment/customfunctions/KnativeTypeHandler.java

@@ -34,6 +34,7 @@
 import org.kie.kogito.serverless.workflow.parser.ParserContext;
 import org.kie.kogito.serverless.workflow.parser.VariableInfo;
 import org.kie.kogito.serverless.workflow.parser.types.WorkItemTypeHandler;
+import org.kie.kogito.serverless.workflow.suppliers.JsonNodeResultHandlerSupplier;
 import org.kie.kogito.serverless.workflow.suppliers.ParamsRestBodyBuilderSupplier;
 import org.kogito.workitem.rest.RestWorkItemHandler;
 
@@ -107,6 +108,7 @@ private static List<String> getPayloadFields(FunctionRef functionRef) {
                 context.getContext(), String.class, DEFAULT_REQUEST_TIMEOUT_VALUE);
 
         return node.workParameter(RestWorkItemHandler.BODY_BUILDER, new ParamsRestBodyBuilderSupplier())
+                .workParameter(RestWorkItemHandler.RESULT_HANDLER, new JsonNodeResultHandlerSupplier())
                 .workParameter(RestWorkItemHandler.REQUEST_TIMEOUT_IN_MILLIS, requestTimeout)
                 .metaData(TaskDescriptor.KEY_WORKITEM_TYPE, RestWorkItemHandler.REST_TASK_TYPE)
                 .workName(KnativeWorkItemHandler.NAME);

quarkus/addons/knative/serving/runtime/src/main/java/org/kie/kogito/addons/quarkus/knative/serving/customfunctions/PlainJsonKnativeParamsDecorator.java

@@ -18,17 +18,23 @@
  */
 package org.kie.kogito.addons.quarkus.knative.serving.customfunctions;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import org.kie.kogito.event.cloudevents.utils.CloudEventUtils;
 import org.kie.kogito.internal.process.workitem.KogitoWorkItem;
 import org.kogito.workitem.rest.decorators.PrefixParamsDecorator;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+
 import io.vertx.mutiny.ext.web.client.HttpRequest;
 
 import static org.kie.kogito.addons.quarkus.knative.serving.customfunctions.KnativeWorkItemHandler.CLOUDEVENT_SENT_AS_PLAIN_JSON_ERROR_MESSAGE;
 import static org.kie.kogito.addons.quarkus.knative.serving.customfunctions.KnativeWorkItemHandler.ID;
+import static org.kie.kogito.serverless.workflow.SWFConstants.MODEL_WORKFLOW_VAR;
 
 public final class PlainJsonKnativeParamsDecorator extends PrefixParamsDecorator {
 
@@ -37,12 +43,42 @@ public void decorate(KogitoWorkItem workItem, Map<String, Object> parameters, Ht
         if (isCloudEvent(KnativeFunctionPayloadSupplier.getPayload(parameters))) {
             throw new IllegalArgumentException(CLOUDEVENT_SENT_AS_PLAIN_JSON_ERROR_MESSAGE);
         }
-
-        super.decorate(workItem, parameters, request);
+        buildFromParams(workItem, parameters, request);
     }
 
     private static boolean isCloudEvent(Map<String, Object> payload) {
         List<String> cloudEventMissingAttributes = CloudEventUtils.getMissingAttributes(payload);
         return !payload.isEmpty() && (cloudEventMissingAttributes.isEmpty() || (cloudEventMissingAttributes.size() == 1 && cloudEventMissingAttributes.contains(ID)));
     }
+
+    private void buildFromParams(KogitoWorkItem workItem, Map<String, Object> parameters, HttpRequest<?> request) {
+        Object modelVar = parameters.get(MODEL_WORKFLOW_VAR);
+
+        if (!(modelVar instanceof ObjectNode objectNode)) {
+            super.decorate(workItem, parameters, request);
+            return;
+        }
+
+        // Flatten the ObjectNode into a plain map so extractHeadersQueries
+        // can remove header/query keys from it without touching the original
+        Map<String, Object> flatModel = flattenObjectNode(objectNode);
+
+        // Extract headers/queries from the flat map (mutates it and the request)
+        Set<String> extractedKeys = super.extractHeadersQueries(workItem, flatModel, request);
+
+        // Apply the same removals to a copy of the original ObjectNode,
+        // then put the sanitized copy back so downstream sees a clean model
+        ObjectNode sanitizedModel = objectNode.deepCopy();
+        sanitizedModel.remove(extractedKeys);
+        parameters.put(MODEL_WORKFLOW_VAR, sanitizedModel);
+    }
+
+    private Map<String, Object> flattenObjectNode(ObjectNode objectNode) {
+        ObjectMapper mapper = new ObjectMapper();
+        Map<String, Object> result = new HashMap<>();
+        objectNode.properties().forEach(entry ->
+            result.put(entry.getKey(), mapper.convertValue(entry.getValue(), Object.class))
+        );
+        return result;
+    }
 }