feat: run nested async workflows (#9119)

packages/core/workflows-sdk/src/helper/workflow-export.ts

@@ -74,6 +74,7 @@ function createContextualWorkflowRunner<
   ) => {
     if (!executionContainer) {
       const container_ = flow.container as MedusaContainer
+
       if (!container_ || !isPresent(container_?.registrations)) {
         executionContainer = MedusaModule.getLoadedModules().map(
           (mod) => Object.values(mod)[0]
@@ -85,12 +86,13 @@ function createContextualWorkflowRunner<
       flow.container = executionContainer
     }
 
-    const { eventGroupId } = context
+    const { eventGroupId, parentStepIdempotencyKey } = context
 
     attachOnFinishReleaseEvents(events, eventGroupId!, flow, { logOnError })
 
     const flowMetadata = {
       eventGroupId,
+      parentStepIdempotencyKey,
     }
 
     const args = [

packages/core/workflows-sdk/src/medusa-workflow.ts

@@ -22,8 +22,8 @@ class MedusaWorkflow {
     MedusaWorkflow.workflows[workflowId] = exportedWorkflow
   }
 
-  static getWorkflow(workflowId) {
-    return MedusaWorkflow.workflows[workflowId]
+  static getWorkflow(workflowId): ExportedWorkflow {
+    return MedusaWorkflow.workflows[workflowId] as unknown as ExportedWorkflow
   }
 }
 

packages/core/workflows-sdk/src/utils/composer/__tests__/index.spec.ts

@@ -240,8 +240,10 @@ describe("Workflow composer", () => {
 
       expect(result).toEqual({ result: "hi from outside" })
 
-      expect(parentContext.transactionId).toEqual("transactionId")
-      expect(parentContext.transactionId).toEqual(childContext.transactionId)
+      expect(parentContext.transactionId).toEqual(expect.any(String))
+      expect(parentContext.transactionId).not.toEqual(
+        childContext.transactionId
+      )
 
       expect(parentContext.eventGroupId).toEqual("eventGroupId")
       expect(parentContext.eventGroupId).toEqual(childContext.eventGroupId)
@@ -287,7 +289,9 @@ describe("Workflow composer", () => {
       expect(result).toEqual({ result: "hi from outside" })
 
       expect(parentContext.transactionId).toBeTruthy()
-      expect(parentContext.transactionId).toEqual(childContext.transactionId)
+      expect(parentContext.transactionId).not.toEqual(
+        childContext.transactionId
+      )
 
       expect(parentContext.eventGroupId).toBeTruthy()
       expect(parentContext.eventGroupId).toEqual(childContext.eventGroupId)

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

@@ -7,6 +7,7 @@ import {
 import { OrchestrationUtils, isString } from "@medusajs/utils"
 import { ulid } from "ulid"
 import { StepResponse, resolveValue } from "./helpers"
+import { createStepHandler } from "./helpers/create-step-handler"
 import { proxify } from "./helpers/proxy"
 import {
   CreateWorkflowComposerContext,
@@ -15,7 +16,6 @@ import {
   StepFunctionResult,
   WorkflowData,
 } from "./type"
-import { createStepHandler } from "./helpers/create-step-handler"
 
 /**
  * The type of invocation function passed to a step.
@@ -65,6 +65,11 @@ export type CompensateFn<T> = (
   context: StepExecutionContext
 ) => unknown | Promise<unknown>
 
+export type LocalStepConfig = { name?: string } & Omit<
+  TransactionStepsDefinition,
+  "next" | "uuid" | "action"
+>
+
 export interface ApplyStepOptions<
   TStepInputs extends {
     [K in keyof TInvokeInput]: WorkflowData<TInvokeInput[K]>
@@ -136,6 +141,8 @@ export function applyStep<
 
     this.flow.addAction(stepName, stepConfig)
 
+    this.isAsync ||= !!(stepConfig.async || stepConfig.compensateAsync)
+
     if (!this.handlers.has(stepName)) {
       this.handlers.set(stepName, handler)
     }
@@ -143,12 +150,7 @@ export function applyStep<
     const ret = {
       __type: OrchestrationUtils.SymbolWorkflowStep,
       __step__: stepName,
-      config: (
-        localConfig: { name?: string } & Omit<
-          TransactionStepsDefinition,
-          "next" | "uuid" | "action"
-        >
-      ) => {
+      config: (localConfig: LocalStepConfig) => {
         const newStepName = localConfig.name ?? stepName
         const newConfig = {
           ...stepConfig,
@@ -160,6 +162,7 @@ export function applyStep<
         this.handlers.set(newStepName, handler)
 
         this.flow.replaceAction(stepConfig.uuid!, newStepName, newConfig)
+        this.isAsync ||= !!(newConfig.async || newConfig.compensateAsync)
 
         ret.__step__ = newStepName
         WorkflowManager.update(this.workflowId, this.flow, this.handlers)
@@ -178,7 +181,7 @@ export function applyStep<
           flagSteps.push(confRef)
         }
 
-        return confRef
+        return confRef as StepFunction<TInvokeInput, TInvokeResultOutput>
       },
       if: (
         input: any,

packages/core/workflows-sdk/src/utils/composer/create-workflow.ts

@@ -5,6 +5,7 @@ import {
 } from "@medusajs/orchestration"
 import { LoadedModule, MedusaContainer } from "@medusajs/types"
 import { OrchestrationUtils, isString } from "@medusajs/utils"
+import { ulid } from "ulid"
 import { exportWorkflow } from "../../helper"
 import { createStep } from "./create-step"
 import { proxify } from "./helpers/proxy"
@@ -104,6 +105,7 @@ export function createWorkflow<TData, TResult, THooks extends any[]>(
     __type: OrchestrationUtils.SymbolMedusaWorkflowComposerContext,
     workflowId: name,
     flow: WorkflowManager.getEmptyTransactionDefinition(),
+    isAsync: false,
     handlers,
     hooks_: {
       declared: [],
@@ -176,21 +178,32 @@ export function createWorkflow<TData, TResult, THooks extends any[]>(
   }: {
     input: TData
   }): ReturnType<StepFunction<TData, TResult>> => {
-    // TODO: Async sub workflow is not supported yet
-    // Info: Once the export workflow can fire the execution through the engine if loaded, the async workflow can be executed,
-    // the step would inherit the async configuration and subscribe to the onFinish event of the sub worklow and mark itself as success or failure
-    return createStep(
-      `${name}-as-step`,
+    const step = createStep(
+      {
+        name: `${name}-as-step`,
+        async: context.isAsync,
+        nested: context.isAsync, // if async we flag this is a nested transaction
+      },
       async (stepInput: TData, stepContext) => {
         const { container, ...sharedContext } = stepContext
 
         const transaction = await workflow.run({
           input: stepInput as any,
           container,
-          context: sharedContext,
+          context: {
+            ...sharedContext,
+            parentStepIdempotencyKey: stepContext.idempotencyKey,
+            transactionId: ulid(),
+          },
         })
 
-        return new StepResponse(transaction.result, transaction)
+        const { result, transaction: flowTransaction } = transaction
+
+        if (!context.isAsync || flowTransaction.hasFinished()) {
+          return new StepResponse(result, transaction)
+        }
+
+        return
       },
       async (transaction, { container }) => {
         if (!transaction) {
@@ -200,6 +213,8 @@ export function createWorkflow<TData, TResult, THooks extends any[]>(
         await workflow(container).cancel(transaction)
       }
     )(input) as ReturnType<StepFunction<TData, TResult>>
+
+    return step
   }
 
   return mainFlow as ReturnWorkflow<TData, TResult, THooks>

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

@@ -1,13 +1,13 @@
+import { WorkflowStepHandlerArguments } from "@medusajs/orchestration"
+import { OrchestrationUtils, deepCopy } from "@medusajs/utils"
+import { ApplyStepOptions } from "../create-step"
 import {
   CreateWorkflowComposerContext,
   StepExecutionContext,
   WorkflowData,
 } from "../type"
-import { WorkflowStepHandlerArguments } from "@medusajs/orchestration"
 import { resolveValue } from "./resolve-value"
 import { StepResponse } from "./step-response"
-import { deepCopy, OrchestrationUtils } from "@medusajs/utils"
-import { ApplyStepOptions } from "../create-step"
 
 export function createStepHandler<
   TInvokeInput,
@@ -36,6 +36,8 @@ export function createStepHandler<
       const idempotencyKey = metadata.idempotency_key
 
       stepArguments.context!.idempotencyKey = idempotencyKey
+
+      const flowMetadata = stepArguments.transaction.getFlow()?.metadata
       const executionContext: StepExecutionContext = {
         workflowId: metadata.model_id,
         stepName: metadata.action,
@@ -45,8 +47,9 @@ export function createStepHandler<
         container: stepArguments.container,
         metadata,
         eventGroupId:
-          stepArguments.transaction.getFlow()?.metadata?.eventGroupId ??
-          stepArguments.context!.eventGroupId,
+          flowMetadata?.eventGroupId ?? stepArguments.context!.eventGroupId,
+        parentStepIdempotencyKey:
+          flowMetadata?.parentStepIdempotencyKey as string,
         transactionId: stepArguments.context!.transactionId,
         context: stepArguments.context!,
       }
@@ -74,11 +77,14 @@ export function createStepHandler<
 
           stepArguments.context!.idempotencyKey = idempotencyKey
 
+          const flowMetadata = stepArguments.transaction.getFlow()?.metadata
           const executionContext: StepExecutionContext = {
             workflowId: metadata.model_id,
             stepName: metadata.action,
             action: "compensate",
             idempotencyKey,
+            parentStepIdempotencyKey:
+              flowMetadata?.parentStepIdempotencyKey as string,
             attempt: metadata.attempt,
             container: stepArguments.container,
             metadata,

packages/core/workflows-sdk/src/utils/composer/type.ts

@@ -98,6 +98,7 @@ export type CreateWorkflowComposerContext = {
   hooksCallback_: Record<string, HookHandler>
   workflowId: string
   flow: OrchestratorBuilder
+  isAsync: boolean
   handlers: WorkflowHandler
   stepBinder: <TOutput = unknown>(
     fn: StepFunctionResult
@@ -127,6 +128,11 @@ export interface StepExecutionContext {
    */
   idempotencyKey: string
 
+  /**
+   * The idempoency key of the parent step.
+   */
+  parentStepIdempotencyKey?: string
+
   /**
    * The name of the step.
    */