feat(workflows-sdk,orchestration): async step as background task (#6886)

packages/workflows-sdk/src/utils/composer/create-step.ts

@@ -1,6 +1,7 @@
 import {
   TransactionStepsDefinition,
   WorkflowManager,
+  WorkflowStepHandler,
   WorkflowStepHandlerArguments,
 } from "@medusajs/orchestration"
 import { OrchestrationUtils, deepCopy, isString } from "@medusajs/utils"
@@ -121,25 +122,23 @@ function applyStep<
     }
 
     const handler = {
-      invoke: async (transactionContext: WorkflowStepHandlerArguments) => {
-        const metadata = transactionContext.metadata
+      invoke: async (stepArguments: WorkflowStepHandlerArguments) => {
+        const metadata = stepArguments.metadata
         const idempotencyKey = metadata.idempotency_key
 
-        transactionContext.context!.idempotencyKey = idempotencyKey
+        stepArguments.context!.idempotencyKey = idempotencyKey
         const executionContext: StepExecutionContext = {
           workflowId: metadata.model_id,
           stepName: metadata.action,
           action: "invoke",
           idempotencyKey,
           attempt: metadata.attempt,
-          container: transactionContext.container,
+          container: stepArguments.container,
           metadata,
-          context: transactionContext.context!,
+          context: stepArguments.context!,
         }
 
-        const argInput = input
-          ? await resolveValue(input, transactionContext)
-          : {}
+        const argInput = input ? await resolveValue(input, stepArguments) : {}
         const stepResponse: StepResponse<any, any> = await invokeFn.apply(
           this,
           [argInput, executionContext]
@@ -156,11 +155,11 @@ function applyStep<
         }
       },
       compensate: compensateFn
-        ? async (transactionContext: WorkflowStepHandlerArguments) => {
-            const metadata = transactionContext.metadata
+        ? async (stepArguments: WorkflowStepHandlerArguments) => {
+            const metadata = stepArguments.metadata
             const idempotencyKey = metadata.idempotency_key
 
-            transactionContext.context!.idempotencyKey = idempotencyKey
+            stepArguments.context!.idempotencyKey = idempotencyKey
 
             const executionContext: StepExecutionContext = {
               workflowId: metadata.model_id,
@@ -168,13 +167,12 @@ function applyStep<
               action: "compensate",
               idempotencyKey,
               attempt: metadata.attempt,
-              container: transactionContext.container,
+              container: stepArguments.container,
               metadata,
-              context: transactionContext.context!,
+              context: stepArguments.context!,
             }
 
-            const stepOutput = (transactionContext.invoke[stepName] as any)
-              ?.output
+            const stepOutput = (stepArguments.invoke[stepName] as any)?.output
             const invokeResult =
               stepOutput?.__type ===
               OrchestrationUtils.SymbolWorkflowStepResponse
@@ -191,6 +189,8 @@ function applyStep<
         : undefined,
     }
 
+    wrapAsyncHandler(stepConfig, handler)
+
     stepConfig.uuid = ulid()
     stepConfig.noCompensation = !compensateFn
 
@@ -231,6 +231,62 @@ function applyStep<
   }
 }
 
+/**
+ * @internal
+ *
+ * Internal function to handle async steps to be automatically marked as completed after they are executed.
+ *
+ * @param stepConfig
+ * @param handle
+ */
+function wrapAsyncHandler(
+  stepConfig: TransactionStepsDefinition,
+  handle: {
+    invoke: WorkflowStepHandler
+    compensate?: WorkflowStepHandler
+  }
+) {
+  if (stepConfig.async) {
+    if (typeof handle.invoke === "function") {
+      const originalInvoke = handle.invoke
+
+      handle.invoke = async (stepArguments: WorkflowStepHandlerArguments) => {
+        const response = (await originalInvoke(stepArguments)) as any
+        if (
+          response?.output?.__type !==
+          OrchestrationUtils.SymbolWorkflowStepResponse
+        ) {
+          return
+        }
+
+        stepArguments.step.definition.backgroundExecution = true
+        return response
+      }
+    }
+  }
+
+  if (stepConfig.compensateAsync) {
+    if (typeof handle.compensate === "function") {
+      const originalCompensate = handle.compensate!
+      handle.compensate = async (
+        stepArguments: WorkflowStepHandlerArguments
+      ) => {
+        const response = (await originalCompensate(stepArguments)) as any
+
+        if (
+          response?.output?.__type !==
+          OrchestrationUtils.SymbolWorkflowStepResponse
+        ) {
+          return
+        }
+        stepArguments.step.definition.backgroundExecution = true
+
+        return response
+      }
+    }
+  }
+}
+
 /**
  * This function creates a {@link StepFunction} that can be used as a step in a workflow constructed by the {@link createWorkflow} function.
  *