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/orchestration/src/workflow/local-workflow.ts

@@ -3,6 +3,8 @@ import { createContainerLike, createMedusaContainer } from "@medusajs/utils"
 import { asValue } from "awilix"
 import {
   DistributedTransaction,
+  DistributedTransactionEvent,
+  DistributedTransactionEvents,
   TransactionOrchestrator,
   TransactionStepsDefinition,
 } from "../transaction"
@@ -79,7 +81,174 @@ export class LocalWorkflow {
     return this.workflow.flow_
   }
 
-  async run(uniqueTransactionId: string, input?: unknown, context?: Context) {
+  private registerEventCallbacks({
+    orchestrator,
+    transaction,
+    subscribe,
+    idempotencyKey,
+  }: {
+    orchestrator: TransactionOrchestrator
+    transaction?: DistributedTransaction
+    subscribe?: DistributedTransactionEvents
+    idempotencyKey?: string
+  }) {
+    const modelId = orchestrator.id
+    let transactionId
+
+    if (transaction) {
+      transactionId = transaction!.transactionId
+    } else if (idempotencyKey) {
+      const [, trxId] = idempotencyKey!.split(":")
+      transactionId = trxId
+    }
+
+    const eventWrapperMap = new Map()
+    for (const [key, handler] of Object.entries(subscribe ?? {})) {
+      eventWrapperMap.set(key, (args) => {
+        const { transaction } = args
+
+        if (
+          transaction.transactionId !== transactionId ||
+          transaction.modelId !== modelId
+        ) {
+          return
+        }
+
+        handler(args)
+      })
+    }
+
+    if (subscribe?.onBegin) {
+      orchestrator.on(
+        DistributedTransactionEvent.BEGIN,
+        eventWrapperMap.get("onBegin")
+      )
+    }
+
+    if (subscribe?.onResume) {
+      orchestrator.on(
+        DistributedTransactionEvent.RESUME,
+        eventWrapperMap.get("onResume")
+      )
+    }
+
+    if (subscribe?.onCompensateBegin) {
+      orchestrator.on(
+        DistributedTransactionEvent.COMPENSATE_BEGIN,
+        eventWrapperMap.get("onCompensateBegin")
+      )
+    }
+
+    if (subscribe?.onTimeout) {
+      orchestrator.on(
+        DistributedTransactionEvent.TIMEOUT,
+        eventWrapperMap.get("onTimeout")
+      )
+    }
+
+    if (subscribe?.onFinish) {
+      orchestrator.on(
+        DistributedTransactionEvent.FINISH,
+        eventWrapperMap.get("onFinish")
+      )
+    }
+
+    const resumeWrapper = ({ transaction }) => {
+      if (
+        transaction.modelId !== modelId ||
+        transaction.transactionId !== transactionId
+      ) {
+        return
+      }
+
+      if (subscribe?.onStepBegin) {
+        transaction.on(
+          DistributedTransactionEvent.STEP_BEGIN,
+          eventWrapperMap.get("onStepBegin")
+        )
+      }
+
+      if (subscribe?.onStepSuccess) {
+        transaction.on(
+          DistributedTransactionEvent.STEP_SUCCESS,
+          eventWrapperMap.get("onStepSuccess")
+        )
+      }
+
+      if (subscribe?.onStepFailure) {
+        transaction.on(
+          DistributedTransactionEvent.STEP_FAILURE,
+          eventWrapperMap.get("onStepFailure")
+        )
+      }
+
+      if (subscribe?.onCompensateStepSuccess) {
+        transaction.on(
+          DistributedTransactionEvent.COMPENSATE_STEP_SUCCESS,
+          eventWrapperMap.get("onCompensateStepSuccess")
+        )
+      }
+
+      if (subscribe?.onCompensateStepFailure) {
+        transaction.on(
+          DistributedTransactionEvent.COMPENSATE_STEP_FAILURE,
+          eventWrapperMap.get("onCompensateStepFailure")
+        )
+      }
+    }
+
+    if (transaction) {
+      resumeWrapper({ transaction })
+    } else {
+      orchestrator.once("resume", resumeWrapper)
+    }
+
+    const cleanUp = () => {
+      subscribe?.onFinish &&
+        orchestrator.removeListener(
+          DistributedTransactionEvent.FINISH,
+          eventWrapperMap.get("onFinish")
+        )
+      subscribe?.onResume &&
+        orchestrator.removeListener(
+          DistributedTransactionEvent.RESUME,
+          eventWrapperMap.get("onResume")
+        )
+      subscribe?.onBegin &&
+        orchestrator.removeListener(
+          DistributedTransactionEvent.BEGIN,
+          eventWrapperMap.get("onBegin")
+        )
+      subscribe?.onCompensateBegin &&
+        orchestrator.removeListener(
+          DistributedTransactionEvent.COMPENSATE_BEGIN,
+          eventWrapperMap.get("onCompensateBegin")
+        )
+      subscribe?.onTimeout &&
+        orchestrator.removeListener(
+          DistributedTransactionEvent.TIMEOUT,
+          eventWrapperMap.get("onTimeout")
+        )
+
+      orchestrator.removeListener(
+        DistributedTransactionEvent.RESUME,
+        resumeWrapper
+      )
+
+      eventWrapperMap.clear()
+    }
+
+    return {
+      cleanUpEventListeners: cleanUp,
+    }
+  }
+
+  async run(
+    uniqueTransactionId: string,
+    input?: unknown,
+    context?: Context,
+    subscribe?: DistributedTransactionEvents
+  ) {
     if (this.flow.hasChanges) {
       this.commit()
     }
@@ -92,36 +261,104 @@ export class LocalWorkflow {
       input
     )
 
+    const { cleanUpEventListeners } = this.registerEventCallbacks({
+      orchestrator,
+      transaction,
+      subscribe,
+    })
+
     await orchestrator.resume(transaction)
 
+    cleanUpEventListeners()
+
+    return transaction
+  }
+
+  async getRunningTransaction(uniqueTransactionId: string, context?: Context) {
+    const { handler, orchestrator } = this.workflow
+
+    const transaction = await orchestrator.retrieveExistingTransaction(
+      uniqueTransactionId,
+      handler(this.container, context)
+    )
+
+    return transaction
+  }
+
+  async cancel(
+    uniqueTransactionId: string,
+    context?: Context,
+    subscribe?: DistributedTransactionEvents
+  ) {
+    const { orchestrator } = this.workflow
+
+    const transaction = await this.getRunningTransaction(
+      uniqueTransactionId,
+      context
+    )
+
+    const { cleanUpEventListeners } = this.registerEventCallbacks({
+      orchestrator,
+      transaction,
+      subscribe,
+    })
+
+    await orchestrator.cancelTransaction(transaction)
+
+    cleanUpEventListeners()
+
     return transaction
   }
 
   async registerStepSuccess(
     idempotencyKey: string,
     response?: unknown,
-    context?: Context
+    context?: Context,
+    subscribe?: DistributedTransactionEvents
   ): Promise<DistributedTransaction> {
     const { handler, orchestrator } = this.workflow
-    return await orchestrator.registerStepSuccess(
+
+    const { cleanUpEventListeners } = this.registerEventCallbacks({
+      orchestrator,
+      idempotencyKey,
+      subscribe,
+    })
+
+    const transaction = await orchestrator.registerStepSuccess(
       idempotencyKey,
       handler(this.container, context),
       undefined,
       response
     )
+
+    cleanUpEventListeners()
+
+    return transaction
   }
 
   async registerStepFailure(
     idempotencyKey: string,
     error?: Error | any,
-    context?: Context
+    context?: Context,
+    subscribe?: DistributedTransactionEvents
   ): Promise<DistributedTransaction> {
     const { handler, orchestrator } = this.workflow
-    return await orchestrator.registerStepFailure(
+
+    const { cleanUpEventListeners } = this.registerEventCallbacks({
+      orchestrator,
+      idempotencyKey,
+      subscribe,
+    })
+
+    const transaction = await orchestrator.registerStepFailure(
       idempotencyKey,
       error,
       handler(this.container, context)
     )
+
+    cleanUpEventListeners()
+
+    return transaction
   }
 
   addAction(