diff --git a/experimental/fluent/func/src/main/java/io/serverlessworkflow/fluent/func/dsl/Step.java b/experimental/fluent/func/src/main/java/io/serverlessworkflow/fluent/func/dsl/Step.java index 2f5055b09..9a1e18626 100644 --- a/experimental/fluent/func/src/main/java/io/serverlessworkflow/fluent/func/dsl/Step.java +++ b/experimental/fluent/func/src/main/java/io/serverlessworkflow/fluent/func/dsl/Step.java @@ -156,12 +156,12 @@ public SELF exportAs(Function function, Class taskResultClass) { /** * Shapes what the task exports for downstream consumers using a context-aware filter function. * - *

This variant provides access to both workflow and task context, allowing you to inspect - * metadata when shaping the export. + *

This variant provides access to the task result, workflow context, and task context, + * allowing you to inspect metadata when shaping the export. * - * @param the workflow context type + * @param the task result type * @param the export type (what gets written in the context) - * @param function the filter function with workflow and task context + * @param function the filter function with task result, workflow context, and task context * @return this step for method chaining * @see io.serverlessworkflow.fluent.func.spi.FuncTaskTransformations#exportAs(FilterFunction) */ @@ -177,10 +177,10 @@ public SELF exportAs(FilterFunction function) { * Shapes what the task exports for downstream consumers using a context-aware filter function * with explicit input type. * - * @param the workflow context type + * @param the task result type * @param the export type (what gets written in the context) - * @param function the filter function with workflow and task context - * @param taskResultClass the class of the workflow context type + * @param function the filter function with task result, workflow context, and task context + * @param taskResultClass the class of the task result type * @return this step for method chaining * @see io.serverlessworkflow.fluent.func.spi.FuncTaskTransformations#exportAs(FilterFunction, * Class) @@ -278,7 +278,7 @@ public SELF exportAs(String jqExpression) { * *

{@code
    * agent("investmentAnalyst", analyst::analyse, InvestmentMemo.class)
-   *     .outputAs(memo -> Map.of("memo", memo))
+   *     .outputAs(taskOutput -> Map.of("memo", taskOutput))
    *     .when(condition);
    * }
* @@ -318,20 +318,21 @@ public SELF outputAs(Function function, Class taskResultClass) { * Shapes what gets written back into the workflow data document using a context-aware filter * function. * - *

This variant provides access to both workflow and task context, allowing you to inspect - * metadata, task input, and raw output when shaping the committed output. + *

This variant provides access to the task result as the first lambda parameter, plus the + * workflow and task context for metadata, task input, and raw output when shaping the committed + * output. * *

Example: * *

{@code
    * get("fetchMarketData", "http://localhost:8081/market-data/{ticker}")
-   *     .outputAs((MarketDataSnapshot snapshot,
+   *     .outputAs((MarketDataSnapshot taskOutput,
    *                WorkflowContextData wf,
    *                TaskContextData task) -> {
    *         var input = task.input().asMap().orElseThrow();
    *         var rawBody = task.rawOutput().asText().orElseThrow();
    *         return new InvestmentPrompt(
-   *             snapshot.ticker(),
+   *             taskOutput.ticker(),
    *             input.get("objective").toString(),
    *             rawBody
    *         );
@@ -340,7 +341,7 @@ public  SELF outputAs(Function function, Class taskResultClass) {
    *
    * @param  the task result type
    * @param  the output type (what gets written to workflow data)
-   * @param function the filter function with workflow and task context
+   * @param function the filter function with task result, workflow context, and task context
    * @return this step for method chaining
    * @see io.serverlessworkflow.fluent.func.spi.FuncTransformations#outputAs(FilterFunction)
    */
@@ -358,7 +359,7 @@ public  SELF outputAs(FilterFunction function) {
    *
    * @param  the task result type
    * @param  the output type (what gets written to workflow data)
-   * @param function the filter function with workflow and task context
+   * @param function the filter function with task result, workflow context, and task context
    * @param taskResultClass the class of the task result type
    * @return this step for method chaining
    * @see io.serverlessworkflow.fluent.func.spi.FuncTransformations#outputAs(FilterFunction, Class)
diff --git a/experimental/lambda/src/test/java/io/serverless/workflow/impl/executors/func/FuncDSLDataFlowTransformationHelpersTest.java b/experimental/lambda/src/test/java/io/serverless/workflow/impl/executors/func/FuncDSLDataFlowTransformationHelpersTest.java
index 7510eada9..ba0896fc0 100644
--- a/experimental/lambda/src/test/java/io/serverless/workflow/impl/executors/func/FuncDSLDataFlowTransformationHelpersTest.java
+++ b/experimental/lambda/src/test/java/io/serverless/workflow/impl/executors/func/FuncDSLDataFlowTransformationHelpersTest.java
@@ -17,7 +17,6 @@
 
 import static io.serverlessworkflow.fluent.func.dsl.FuncDSL.function;
 import static io.serverlessworkflow.fluent.func.dsl.FuncDSL.input;
-import static io.serverlessworkflow.fluent.func.dsl.FuncDSL.output;
 
 import io.serverlessworkflow.api.types.Workflow;
 import io.serverlessworkflow.fluent.func.FuncWorkflowBuilder;
@@ -82,11 +81,13 @@ void test_input_with_outputAs() {
                         },
                         Long.class)
                     .outputAs(
-                        (object, workflowContext, taskContextData) -> {
-                          softly.assertThat(object).isEqualTo(15L);
+                        (Long taskOutput,
+                            WorkflowContextData workflowContext,
+                            TaskContextData taskContextData) -> {
+                          softly.assertThat(taskOutput).isEqualTo(15L);
                           Long input = input(workflowContext, Long.class);
                           softly.assertThat(input).isEqualTo(10L);
-                          return input + object;
+                          return input + taskOutput;
                         },
                         Long.class))
             .build();
@@ -119,10 +120,9 @@ void test_output_with_exportAs() {
                         },
                         Long.class)
                     .exportAs(
-                        (Long object,
+                        (Long taskOutput,
                             WorkflowContextData workflowContext,
                             TaskContextData taskContextData) -> {
-                          Long taskOutput = output(taskContextData, Long.class);
                           softly.assertThat(taskOutput).isEqualTo(15L);
                           return taskOutput * 2;
                         }))
@@ -138,6 +138,53 @@ void test_output_with_exportAs() {
     softly.assertAll();
   }
 
+  @Test
+  void test_output_with_outputAs_and_exportAs() {
+
+    SoftAssertions softly = new SoftAssertions();
+
+    Workflow workflow =
+        FuncWorkflowBuilder.workflow("enrichOutputWithOutputAsAndExportAsTest")
+            .tasks(
+                function(
+                        "add5",
+                        (Long input) -> {
+                          softly.assertThat(input).isEqualTo(10L);
+                          return input + 5;
+                        },
+                        Long.class)
+                    .outputAs(
+                        (Long taskOutput) -> {
+                          softly.assertThat(taskOutput).isEqualTo(15L);
+                          return taskOutput * 2;
+                        })
+                    .exportAs(
+                        (Long taskOutput,
+                            WorkflowContextData workflowContext,
+                            TaskContextData taskContextData) -> {
+                          softly.assertThat(taskOutput).isEqualTo(30L);
+                          softly
+                              .assertThat(taskContextData.output().asNumber().orElseThrow())
+                              .isEqualTo(30L);
+                          return taskOutput;
+                        }))
+            .build();
+
+    try (WorkflowApplication app = WorkflowApplication.builder().build()) {
+      WorkflowDefinition def = app.workflowDefinition(workflow);
+      WorkflowModel model = def.instance(10L).start().join();
+      Number number = model.asNumber().orElseThrow();
+      softly.assertThat(number.longValue()).isEqualTo(30L);
+
+      WorkflowInstance instance = def.instance(10L);
+      instance.start().join();
+      Number contextNumber = instance.context().asNumber().orElseThrow();
+      softly.assertThat(contextNumber.longValue()).isEqualTo(30L);
+    }
+
+    softly.assertAll();
+  }
+
   @Test
   void test_input_with_inputFrom_fluent_way() {
     SoftAssertions softly = new SoftAssertions();
@@ -178,10 +225,9 @@ void test_input_with_exportAs() {
                         },
                         Long.class)
                     .exportAs(
-                        (Long object,
+                        (Long taskOutput,
                             WorkflowContextData workflowContext,
                             TaskContextData taskContextData) -> {
-                          Long taskOutput = output(taskContextData, Long.class);
                           softly.assertThat(taskOutput).isEqualTo(15L);
                           Long input = input(workflowContext, Long.class);
                           softly.assertThat(input).isEqualTo(10L);
diff --git a/impl/core/src/main/java/io/serverlessworkflow/impl/WorkflowUtils.java b/impl/core/src/main/java/io/serverlessworkflow/impl/WorkflowUtils.java
index 16979fc1f..011fb2fa7 100644
--- a/impl/core/src/main/java/io/serverlessworkflow/impl/WorkflowUtils.java
+++ b/impl/core/src/main/java/io/serverlessworkflow/impl/WorkflowUtils.java
@@ -82,11 +82,16 @@ public static Optional buildWorkflowFilter(WorkflowApplication a
   }
 
   public static Optional buildWorkflowFilter(WorkflowApplication app, ExportAs as) {
-    return as != null
-        ? Optional.of(
-            buildFilterFromStrObject(
-                app.expressionFactory(), app.contextFactory(), as.getString(), as.getObject()))
-        : Optional.empty();
+    if (as == null) {
+      return Optional.empty();
+    }
+    WorkflowFilter filter =
+        buildFilterFromStrObject(
+            app.expressionFactory(), app.contextFactory(), as.getString(), as.getObject());
+    if (as.getString() == null && as.getObject() != null) {
+      return Optional.of((w, t, m) -> filter.apply(w, t, t.output()));
+    }
+    return Optional.of(filter);
   }
 
   public static WorkflowFilter buildWorkflowFilter(WorkflowApplication app, Object obj) {