feat(orchestration, workflows-sdk): Add events management and implementation to manage async workflows (#5951)

* feat(orchestration, workflows-sdk): Add events management and implementation to manage async workflows

* Create fresh-boxes-scream.md

* cleanup

* fix: resolveValue input ref

* resolve value recursive

* chore: resolve result value only for new api

* chore: save checkpoint before scheduling

* features

* fix: beginTransaction checking existing transaction

---------

Co-authored-by: Carlos R. L. Rodrigues <rodrigolr@gmail.com>
Co-authored-by: Carlos R. L. Rodrigues <37986729+carlos-r-l-rodrigues@users.noreply.github.com>

packages/workflows-sdk/src/helper/__tests__/workflow-export.spec.ts

@@ -23,6 +23,28 @@ jest.mock("@medusajs/orchestration", () => {
             }),
           }
         }),
+        registerStepSuccess: jest.fn(() => {
+          return {
+            getErrors: jest.fn(),
+            getState: jest.fn(() => "done"),
+            getContext: jest.fn(() => {
+              return {
+                invoke: { result_step: "invoke_test" },
+              }
+            }),
+          }
+        }),
+        registerStepFailure: jest.fn(() => {
+          return {
+            getErrors: jest.fn(),
+            getState: jest.fn(() => "done"),
+            getContext: jest.fn(() => {
+              return {
+                invoke: { result_step: "invoke_test" },
+              }
+            }),
+          }
+        }),
       }
     }),
   }

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

@@ -1,5 +1,6 @@
 import {
   DistributedTransaction,
+  DistributedTransactionEvents,
   LocalWorkflow,
   TransactionHandlerType,
   TransactionState,
@@ -8,15 +9,36 @@ import {
 import { Context, LoadedModule, MedusaContainer } from "@medusajs/types"
 
 import { MedusaModule } from "@medusajs/modules-sdk"
+import { OrchestrationUtils } from "@medusajs/utils"
 import { EOL } from "os"
 import { ulid } from "ulid"
-import { OrchestrationUtils } from "@medusajs/utils"
+import { MedusaWorkflow } from "../medusa-workflow"
+import { resolveValue } from "../utils/composer"
 
 export type FlowRunOptions<TData = unknown> = {
   input?: TData
   context?: Context
-  resultFrom?: string | string[]
+  resultFrom?: string | string[] | Symbol
   throwOnError?: boolean
+  events?: DistributedTransactionEvents
+}
+
+export type FlowRegisterStepSuccessOptions<TData = unknown> = {
+  idempotencyKey: string
+  response?: TData
+  context?: Context
+  resultFrom?: string | string[] | Symbol
+  throwOnError?: boolean
+  events?: DistributedTransactionEvents
+}
+
+export type FlowRegisterStepFailureOptions<TData = unknown> = {
+  idempotencyKey: string
+  response?: TData
+  context?: Context
+  resultFrom?: string | string[] | Symbol
+  throwOnError?: boolean
+  events?: DistributedTransactionEvents
 }
 
 export type WorkflowResult<TResult = unknown> = {
@@ -25,24 +47,59 @@ export type WorkflowResult<TResult = unknown> = {
   result: TResult
 }
 
+export type ExportedWorkflow<
+  TData = unknown,
+  TResult = unknown,
+  TDataOverride = undefined,
+  TResultOverride = undefined
+> = {
+  run: (
+    args?: FlowRunOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    >
+  ) => Promise<
+    WorkflowResult<
+      TResultOverride extends undefined ? TResult : TResultOverride
+    >
+  >
+  registerStepSuccess: (
+    args?: FlowRegisterStepSuccessOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    >
+  ) => Promise<
+    WorkflowResult<
+      TResultOverride extends undefined ? TResult : TResultOverride
+    >
+  >
+  registerStepFailure: (
+    args?: FlowRegisterStepFailureOptions<
+      TDataOverride extends undefined ? TData : TDataOverride
+    >
+  ) => Promise<
+    WorkflowResult<
+      TResultOverride extends undefined ? TResult : TResultOverride
+    >
+  >
+}
+
 export const exportWorkflow = <TData = unknown, TResult = unknown>(
   workflowId: string,
-  defaultResult?: string,
-  dataPreparation?: (data: TData) => Promise<unknown>
+  defaultResult?: string | Symbol,
+  dataPreparation?: (data: TData) => Promise<unknown>,
+  options?: {
+    wrappedInput?: boolean
+  }
 ) => {
-  return function <TDataOverride = undefined, TResultOverride = undefined>(
+  function exportedWorkflow<
+    TDataOverride = undefined,
+    TResultOverride = undefined
+  >(
     container?: LoadedModule[] | MedusaContainer
-  ): Omit<LocalWorkflow, "run"> & {
-    run: (
-      args?: FlowRunOptions<
-        TDataOverride extends undefined ? TData : TDataOverride
-      >
-    ) => Promise<
-      WorkflowResult<
-        TResultOverride extends undefined ? TResult : TResultOverride
-      >
-    >
-  } {
+  ): Omit<
+    LocalWorkflow,
+    "run" | "registerStepSuccess" | "registerStepFailure"
+  > &
+    ExportedWorkflow<TData, TResult, TDataOverride, TResultOverride> {
     if (!container) {
       container = MedusaModule.getLoadedModules().map(
         (mod) => Object.values(mod)[0]
@@ -52,8 +109,64 @@ export const exportWorkflow = <TData = unknown, TResult = unknown>(
     const flow = new LocalWorkflow(workflowId, container)
 
     const originalRun = flow.run.bind(flow)
+    const originalRegisterStepSuccess = flow.registerStepSuccess.bind(flow)
+    const originalRegisterStepFailure = flow.registerStepFailure.bind(flow)
+
+    const originalExecution = async (
+      method,
+      { throwOnError, resultFrom },
+      ...args
+    ) => {
+      const transaction = await method.apply(method, args)
+
+      const errors = transaction.getErrors(TransactionHandlerType.INVOKE)
+
+      const failedStatus = [TransactionState.FAILED, TransactionState.REVERTED]
+      if (failedStatus.includes(transaction.getState()) && throwOnError) {
+        const errorMessage = errors
+          ?.map((err) => `${err.error?.message}${EOL}${err.error?.stack}`)
+          ?.join(`${EOL}`)
+        throw new Error(errorMessage)
+      }
+
+      let result: any = undefined
+
+      const resFrom =
+        resultFrom?.__type === OrchestrationUtils.SymbolWorkflowStep
+          ? resultFrom.__step__
+          : resultFrom
+
+      if (resFrom) {
+        if (Array.isArray(resFrom)) {
+          result = resFrom.map((from) => {
+            const res = transaction.getContext().invoke?.[from]
+            return res?.__type === OrchestrationUtils.SymbolWorkflowWorkflowData
+              ? res.output
+              : res
+          })
+        } else {
+          const res = transaction.getContext().invoke?.[resFrom]
+          result =
+            res?.__type === OrchestrationUtils.SymbolWorkflowWorkflowData
+              ? res.output
+              : res
+        }
+
+        const ret = result || resFrom
+        result = options?.wrappedInput
+          ? await resolveValue(ret, transaction.getContext())
+          : ret
+      }
+
+      return {
+        errors,
+        transaction,
+        result,
+      }
+    }
+
     const newRun = async (
-      { input, context, throwOnError, resultFrom }: FlowRunOptions = {
+      { input, context, throwOnError, resultFrom, events }: FlowRunOptions = {
         throwOnError: true,
         resultFrom: defaultResult,
       }
@@ -77,59 +190,77 @@ export const exportWorkflow = <TData = unknown, TResult = unknown>(
         }
       }
 
-      const transaction = await originalRun(
+      return await originalExecution(
+        originalRun,
+        { throwOnError, resultFrom },
         context?.transactionId ?? ulid(),
         input,
-        context
+        context,
+        events
       )
-
-      const errors = transaction.getErrors(TransactionHandlerType.INVOKE)
-
-      const failedStatus = [TransactionState.FAILED, TransactionState.REVERTED]
-      if (failedStatus.includes(transaction.getState()) && throwOnError) {
-        const errorMessage = errors
-          ?.map((err) => `${err.error?.message}${EOL}${err.error?.stack}`)
-          ?.join(`${EOL}`)
-        throw new Error(errorMessage)
-      }
-
-      let result: any = undefined
-
-      if (resultFrom) {
-        if (Array.isArray(resultFrom)) {
-          result = resultFrom.map((from) => {
-            const res = transaction.getContext().invoke?.[from]
-            return res?.__type === OrchestrationUtils.SymbolWorkflowWorkflowData
-              ? res.output
-              : res
-          })
-        } else {
-          const res = transaction.getContext().invoke?.[resultFrom]
-          result =
-            res?.__type === OrchestrationUtils.SymbolWorkflowWorkflowData
-              ? res.output
-              : res
-        }
-      }
-
-      return {
-        errors,
-        transaction,
-        result,
-      }
     }
     flow.run = newRun as any
 
-    return flow as unknown as LocalWorkflow & {
-      run: (
-        args?: FlowRunOptions<
-          TDataOverride extends undefined ? TData : TDataOverride
-        >
-      ) => Promise<
-        WorkflowResult<
-          TResultOverride extends undefined ? TResult : TResultOverride
-        >
-      >
+    const newRegisterStepSuccess = async (
+      {
+        response,
+        idempotencyKey,
+        context,
+        throwOnError,
+        resultFrom,
+        events,
+      }: FlowRegisterStepSuccessOptions = {
+        idempotencyKey: "",
+        throwOnError: true,
+        resultFrom: defaultResult,
+      }
+    ) => {
+      resultFrom ??= defaultResult
+      throwOnError ??= true
+
+      return await originalExecution(
+        originalRegisterStepSuccess,
+        { throwOnError, resultFrom },
+        idempotencyKey,
+        response,
+        context,
+        events
+      )
     }
+    flow.registerStepSuccess = newRegisterStepSuccess as any
+
+    const newRegisterStepFailure = async (
+      {
+        response,
+        idempotencyKey,
+        context,
+        throwOnError,
+        resultFrom,
+        events,
+      }: FlowRegisterStepFailureOptions = {
+        idempotencyKey: "",
+        throwOnError: true,
+        resultFrom: defaultResult,
+      }
+    ) => {
+      resultFrom ??= defaultResult
+      throwOnError ??= true
+
+      return await originalExecution(
+        originalRegisterStepFailure,
+        { throwOnError, resultFrom },
+        idempotencyKey,
+        response,
+        context,
+        events
+      )
+    }
+    flow.registerStepFailure = newRegisterStepFailure as any
+
+    return flow as unknown as LocalWorkflow &
+      ExportedWorkflow<TData, TResult, TDataOverride, TResultOverride>
   }
+
+  MedusaWorkflow.registerWorkflow(workflowId, exportedWorkflow)
+  return exportedWorkflow
 }

packages/workflows-sdk/src/index.ts

@@ -1,3 +1,4 @@
 export * from "./helper"
+export * from "./medusa-workflow"
 export * from "./utils/composer"
 export * as Composer from "./utils/composer"

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

@@ -0,0 +1,27 @@
+import { LocalWorkflow } from "@medusajs/orchestration"
+import { LoadedModule, MedusaContainer } from "@medusajs/types"
+import { ExportedWorkflow } from "./helper"
+
+export class MedusaWorkflow {
+  static workflows: Record<
+    string,
+    (
+      container?: LoadedModule[] | MedusaContainer
+    ) => Omit<
+      LocalWorkflow,
+      "run" | "registerStepSuccess" | "registerStepFailure"
+    > &
+      ExportedWorkflow
+  > = {}
+
+  static registerWorkflow(workflowId, exportedWorkflow) {
+    MedusaWorkflow.workflows[workflowId] = exportedWorkflow
+  }
+
+  static getWorkflow(workflowId) {
+    return MedusaWorkflow.workflows[workflowId]
+  }
+}
+
+global.MedusaWorkflow ??= MedusaWorkflow
+exports.MedusaWorkflow = global.MedusaWorkflow

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

@@ -5,8 +5,7 @@ import {
 } from "@medusajs/orchestration"
 import { LoadedModule, MedusaContainer } from "@medusajs/types"
 import { OrchestrationUtils } from "@medusajs/utils"
-import { exportWorkflow, FlowRunOptions, WorkflowResult } from "../../helper"
-import { resolveValue } from "./helpers"
+import { ExportedWorkflow, exportWorkflow } from "../../helper"
 import { proxify } from "./helpers/proxy"
 import {
   CreateWorkflowComposerContext,
@@ -66,17 +65,11 @@ global[OrchestrationUtils.SymbolMedusaWorkflowComposerContext] = null
 type ReturnWorkflow<TData, TResult, THooks extends Record<string, Function>> = {
   <TDataOverride = undefined, TResultOverride = undefined>(
     container?: LoadedModule[] | MedusaContainer
-  ): Omit<LocalWorkflow, "run"> & {
-    run: (
-      args?: FlowRunOptions<
-        TDataOverride extends undefined ? TData : TDataOverride
-      >
-    ) => Promise<
-      WorkflowResult<
-        TResultOverride extends undefined ? TResult : TResultOverride
-      >
-    >
-  }
+  ): Omit<
+    LocalWorkflow,
+    "run" | "registerStepSuccess" | "registerStepFailure"
+  > &
+    ExportedWorkflow<TData, TResult, TDataOverride, TResultOverride>
 } & THooks & {
     getName: () => string
   }
@@ -198,43 +191,19 @@ export function createWorkflow<
 
   WorkflowManager.update(name, context.flow, handlers)
 
-  const workflow = exportWorkflow<TData, TResult>(name)
+  const workflow = exportWorkflow<TData, TResult>(
+    name,
+    returnedStep,
+    undefined,
+    {
+      wrappedInput: true,
+    }
+  )
 
   const mainFlow = <TDataOverride = undefined, TResultOverride = undefined>(
     container?: LoadedModule[] | MedusaContainer
   ) => {
     const workflow_ = workflow<TDataOverride, TResultOverride>(container)
-    const originalRun = workflow_.run
-
-    workflow_.run = (async (
-      args?: FlowRunOptions<
-        TDataOverride extends undefined ? TData : TDataOverride
-      >
-    ): Promise<
-      WorkflowResult<
-        TResultOverride extends undefined ? TResult : TResultOverride
-      >
-    > => {
-      args ??= {}
-      args.resultFrom ??=
-        returnedStep?.__type === OrchestrationUtils.SymbolWorkflowStep
-          ? returnedStep.__step__
-          : undefined
-
-      // Forwards the input to the ref object on composer.apply
-      const workflowResult = (await originalRun(
-        args
-      )) as unknown as WorkflowResult<
-        TResultOverride extends undefined ? TResult : TResultOverride
-      >
-
-      workflowResult.result = await resolveValue(
-        workflowResult.result || returnedStep,
-        workflowResult.transaction.getContext()
-      )
-
-      return workflowResult
-    }) as any
 
     return workflow_
   }

packages/workflows-sdk/src/utils/composer/helpers/resolve-value.ts

@@ -41,7 +41,7 @@ export async function resolveValue(input, transactionContext) {
 
     if (Array.isArray(inputTOUnwrap)) {
       return await promiseAll(
-        inputTOUnwrap.map((i) => unwrapInput(i, transactionContext))
+        inputTOUnwrap.map((i) => resolveValue(i, transactionContext))
       )
     }