feat: Ensure async workflow executions have access to shared container (#8157)

* feat: Ensure async workflow executions have access to shared container

* fix: Register workflow worker on application start

packages/core/medusa-test-utils/src/init-modules.ts

@@ -56,6 +56,8 @@ export async function initModules({
     injectedDependencies,
   })
 
+  await medusaApp.onApplicationStart()
+
   async function shutdown() {
     if (shouldDestroyConnectionAutomatically) {
       await medusaApp.onApplicationPrepareShutdown()

packages/core/modules-sdk/src/definitions.ts

@@ -145,6 +145,7 @@ export const ModulesDefinition: {
     isRequired: false,
     isQueryable: true,
     dependencies: ["logger"],
+    __passSharedContainer: true,
     defaultModuleDeclaration: {
       scope: MODULE_SCOPE.INTERNAL,
       resources: MODULE_RESOURCE_TYPE.SHARED,

packages/core/modules-sdk/src/loaders/utils/load-internal.ts

@@ -139,6 +139,15 @@ export async function loadInternalModule(
     )
   }
 
+  if (resolution.definition.__passSharedContainer) {
+    localContainer.register(
+      "sharedContainer",
+      asFunction(() => {
+        return container
+      })
+    )
+  }
+
   const loaders = moduleResources.loaders ?? loadedModule?.loaders ?? []
   const error = await runLoaders(loaders, {
     container,

packages/core/modules-sdk/src/medusa-app.ts

@@ -233,6 +233,7 @@ export type MedusaAppOutput = {
   generateMigrations: GenerateMigrations
   onApplicationShutdown: () => Promise<void>
   onApplicationPrepareShutdown: () => Promise<void>
+  onApplicationStart: () => Promise<void>
   sharedContainer?: MedusaContainer
 }
 
@@ -284,6 +285,10 @@ async function MedusaApp_({
     await promiseAll([MedusaModule.onApplicationPrepareShutdown()])
   }
 
+  const onApplicationStart = async () => {
+    await MedusaModule.onApplicationStart()
+  }
+
   const modules: MedusaModuleConfig =
     modulesConfig ??
     (
@@ -349,6 +354,7 @@ async function MedusaApp_({
     return {
       onApplicationShutdown,
       onApplicationPrepareShutdown,
+      onApplicationStart,
       modules: allModules,
       link: undefined,
       query: async () => {
@@ -536,6 +542,7 @@ async function MedusaApp_({
   return {
     onApplicationShutdown,
     onApplicationPrepareShutdown,
+    onApplicationStart,
     modules: allModules,
     link: remoteLink,
     query,
@@ -551,11 +558,7 @@ async function MedusaApp_({
 export async function MedusaApp(
   options: MedusaAppOptions = {}
 ): Promise<MedusaAppOutput> {
-  try {
-    return await MedusaApp_(options)
-  } finally {
-    MedusaModule.onApplicationStart(options.onApplicationStartCb)
-  }
+  return await MedusaApp_(options)
 }
 
 export async function MedusaAppMigrateUp(

packages/core/types/src/modules-sdk/index.ts

@@ -87,6 +87,8 @@ export type ModuleDefinition = {
   isRequired?: boolean
   isQueryable?: boolean // If the module is queryable via Remote Joiner
   dependencies?: string[]
+  /** @internal only used in exceptional cases - relying on the shared contrainer breaks encapsulation */
+  __passSharedContainer?: boolean
   defaultModuleDeclaration:
     | InternalModuleDeclaration
     | ExternalModuleDeclaration

packages/medusa/src/loaders/index.ts

@@ -153,14 +153,17 @@ export default async ({
 
   const plugins = getResolvedPlugins(rootDirectory, configModule, true) || []
   const pluginLinks = await resolvePluginsLinks(plugins, container)
+
+  const {
+    onApplicationStart,
+    onApplicationShutdown,
+    onApplicationPrepareShutdown,
+  } = await loadMedusaApp({
+    container,
+    linkModules: pluginLinks,
+  })
+
   await registerWorkflows(plugins)
-
-  const { onApplicationShutdown, onApplicationPrepareShutdown } =
-    await loadMedusaApp({
-      container,
-      linkModules: pluginLinks,
-    })
-
   const entrypointsShutdown = await loadEntrypoints(
     plugins,
     container,
@@ -168,6 +171,7 @@ export default async ({
     rootDirectory
   )
   await createDefaultsWorkflow(container).run()
+  await onApplicationStart()
 
   const shutdown = async () => {
     const pgConnection = container.resolve(

packages/modules/workflow-engine-inmemory/integration-tests/__fixtures__/workflow_scheduled.ts

@@ -6,8 +6,10 @@ import {
 } from "@medusajs/workflows-sdk"
 
 export const createScheduled = (name: string, schedule?: SchedulerOptions) => {
-  const workflowScheduledStepInvoke = jest.fn((input, context) => {
-    return new StepResponse({})
+  const workflowScheduledStepInvoke = jest.fn((input, { container }) => {
+    return new StepResponse({
+      testValue: container.resolve("test-value"),
+    })
   })
 
   const step = createStep("step_1", workflowScheduledStepInvoke)

packages/modules/workflow-engine-inmemory/integration-tests/__tests__/index.spec.ts

@@ -21,6 +21,7 @@ import {
 } from "../__fixtures__/workflow_event_group_id"
 import { createScheduled } from "../__fixtures__/workflow_scheduled"
 import { WorkflowsModuleService } from "@services"
+import { asFunction } from "awilix"
 
 jest.setTimeout(100000)
 
@@ -367,6 +368,25 @@ moduleIntegrationTestRunner<IWorkflowEngineService>({
           )
         })
 
+        it("the scheduled workflow should have access to the shared container", async () => {
+          const sharedContainer =
+            workflowOrcModule["workflowOrchestratorService_"]["container_"]
+
+          sharedContainer.register(
+            "test-value",
+            asFunction(() => "test")
+          )
+
+          const spy = await createScheduled("remove-scheduled", {
+            cron: "* * * * * *",
+          })
+          await jest.runOnlyPendingTimersAsync()
+          expect(spy).toHaveBeenCalledTimes(1)
+          expect(spy).toHaveReturnedWith(
+            expect.objectContaining({ output: { testValue: "test" } })
+          )
+        })
+
         it("should fetch an idempotent workflow after its completion", async () => {
           const { transaction: firstRun } = await workflowOrcModule.run(
             "workflow_idempotent",

packages/modules/workflow-engine-inmemory/src/services/workflow-orchestrator.ts

@@ -80,13 +80,17 @@ const AnySubscriber = "any"
 
 export class WorkflowOrchestratorService {
   private subscribers: Subscribers = new Map()
+  private container_: MedusaContainer
 
   constructor({
     inMemoryDistributedTransactionStorage,
+    sharedContainer,
   }: {
     inMemoryDistributedTransactionStorage: InMemoryDistributedTransactionStorage
     workflowOrchestratorService: WorkflowOrchestratorService
+    sharedContainer: MedusaContainer
   }) {
+    this.container_ = sharedContainer
     inMemoryDistributedTransactionStorage.setWorkflowOrchestratorService(this)
     DistributedTransaction.setStorage(inMemoryDistributedTransactionStorage)
     WorkflowScheduler.setStorage(inMemoryDistributedTransactionStorage)
@@ -136,7 +140,9 @@ export class WorkflowOrchestratorService {
       )
     }
 
-    const flow = exportedWorkflow(container as MedusaContainer)
+    const flow = exportedWorkflow(
+      (container as MedusaContainer) ?? this.container_
+    )
 
     const ret = await flow.run({
       input,
@@ -191,7 +197,9 @@ export class WorkflowOrchestratorService {
       throw new Error(`Workflow with id "${workflowId}" not found.`)
     }
 
-    const flow = exportedWorkflow(container as MedusaContainer)
+    const flow = exportedWorkflow(
+      (container as MedusaContainer) ?? this.container_
+    )
 
     const transaction = await flow.getRunningTransaction(transactionId, context)
 
@@ -227,7 +235,9 @@ export class WorkflowOrchestratorService {
       throw new Error(`Workflow with id "${workflowId}" not found.`)
     }
 
-    const flow = exportedWorkflow(container as MedusaContainer)
+    const flow = exportedWorkflow(
+      (container as MedusaContainer) ?? this.container_
+    )
 
     const events = this.buildWorkflowEvents({
       customEventHandlers: eventHandlers,
@@ -287,7 +297,9 @@ export class WorkflowOrchestratorService {
       throw new Error(`Workflow with id "${workflowId}" not found.`)
     }
 
-    const flow = exportedWorkflow(container as MedusaContainer)
+    const flow = exportedWorkflow(
+      (container as MedusaContainer) ?? this.container_
+    )
 
     const events = this.buildWorkflowEvents({
       customEventHandlers: eventHandlers,

packages/modules/workflow-engine-inmemory/src/services/workflows-module.ts

@@ -2,6 +2,7 @@ import {
   Context,
   DAL,
   InternalModuleDeclaration,
+  MedusaContainer,
   ModulesSdkTypes,
   WorkflowsSdkTypes,
 } from "@medusajs/types"
@@ -31,6 +32,7 @@ export class WorkflowsModuleService<
   protected baseRepository_: DAL.RepositoryService
   protected workflowExecutionService_: ModulesSdkTypes.IMedusaInternalService<TWorkflowExecution>
   protected workflowOrchestratorService_: WorkflowOrchestratorService
+  protected container_: MedusaContainer
 
   constructor(
     {

packages/modules/workflow-engine-redis/integration-tests/__fixtures__/workflow_scheduled.ts

@@ -6,8 +6,10 @@ import {
 } from "@medusajs/workflows-sdk"
 
 export const createScheduled = (name: string, schedule?: SchedulerOptions) => {
-  const workflowScheduledStepInvoke = jest.fn((input, context) => {
-    return new StepResponse({})
+  const workflowScheduledStepInvoke = jest.fn((input, { container }) => {
+    return new StepResponse({
+      testValue: container.resolve("test-value"),
+    })
   })
 
   const step = createStep("step_1", workflowScheduledStepInvoke)

packages/modules/workflow-engine-redis/integration-tests/__tests__/index.spec.ts

@@ -17,7 +17,7 @@ import {
   TransactionHandlerType,
   TransactionStepState,
 } from "@medusajs/utils"
-import { asValue } from "awilix"
+import { asFunction, asValue } from "awilix"
 import { knex } from "knex"
 import { setTimeout } from "timers/promises"
 import "../__fixtures__"
@@ -54,6 +54,7 @@ describe("Workflow Orchestrator module", function () {
       query: remoteQuery,
       modules,
       sharedContainer,
+      onApplicationStart,
     } = await MedusaApp({
       sharedContainer: container,
       sharedResourcesConfig: {
@@ -73,6 +74,8 @@ describe("Workflow Orchestrator module", function () {
       },
     })
 
+    await onApplicationStart()
+
     query = remoteQuery
     sharedContainer_ = sharedContainer!
 
@@ -381,5 +384,21 @@ describe("Workflow Orchestrator module", function () {
         "Tried to execute a scheduled workflow with ID remove-scheduled that does not exist, removing it from the scheduler."
       )
     })
+
+    it("the scheduled workflow should have access to the shared container", async () => {
+      sharedContainer_.register(
+        "test-value",
+        asFunction(() => "test")
+      )
+
+      const spy = await createScheduled("remove-scheduled", {
+        cron: "* * * * * *",
+      })
+      await setTimeout(1100)
+      expect(spy).toHaveBeenCalledTimes(1)
+      expect(spy).toHaveReturnedWith(
+        expect.objectContaining({ output: { testValue: "test" } })
+      )
+    })
   })
 })

packages/modules/workflow-engine-redis/src/services/workflow-orchestrator.ts

@@ -83,6 +83,7 @@ export class WorkflowOrchestratorService {
   private instanceId = ulid()
   protected redisPublisher: Redis
   protected redisSubscriber: Redis
+  protected container_: MedusaContainer
   private subscribers: Subscribers = new Map()
   private activeStepsCount: number = 0
   private logger: Logger
@@ -95,6 +96,7 @@ export class WorkflowOrchestratorService {
     redisPublisher,
     redisSubscriber,
     logger,
+    sharedContainer,
   }: {
     dataLoaderOnly: boolean
     redisDistributedTransactionStorage: RedisDistributedTransactionStorage
@@ -102,7 +104,9 @@ export class WorkflowOrchestratorService {
     redisPublisher: Redis
     redisSubscriber: Redis
     logger: Logger
+    sharedContainer: MedusaContainer
   }) {
+    this.container_ = sharedContainer
     this.redisPublisher = redisPublisher
     this.redisSubscriber = redisSubscriber
     this.logger = logger
@@ -137,6 +141,10 @@ export class WorkflowOrchestratorService {
     }
   }
 
+  async onApplicationStart() {
+    await this.redisDistributedTransactionStorage_.onApplicationStart()
+  }
+
   @InjectSharedContext()
   async run<T = unknown>(
     workflowIdOrWorkflow: string | ReturnWorkflow<any, any, any>,
@@ -175,7 +183,9 @@ export class WorkflowOrchestratorService {
       throw new Error(`Workflow with id "${workflowId}" not found.`)
     }
 
-    const flow = exportedWorkflow(container as MedusaContainer)
+    const flow = exportedWorkflow(
+      (container as MedusaContainer) ?? this.container_
+    )
 
     const ret = await flow.run({
       input,
@@ -230,7 +240,9 @@ export class WorkflowOrchestratorService {
       throw new Error(`Workflow with id "${workflowId}" not found.`)
     }
 
-    const flow = exportedWorkflow(container as MedusaContainer)
+    const flow = exportedWorkflow(
+      (container as MedusaContainer) ?? this.container_
+    )
 
     const transaction = await flow.getRunningTransaction(transactionId, context)
 
@@ -266,7 +278,9 @@ export class WorkflowOrchestratorService {
       throw new Error(`Workflow with id "${workflowId}" not found.`)
     }
 
-    const flow = exportedWorkflow(container as MedusaContainer)
+    const flow = exportedWorkflow(
+      (container as MedusaContainer) ?? this.container_
+    )
 
     const events = this.buildWorkflowEvents({
       customEventHandlers: eventHandlers,
@@ -326,7 +340,9 @@ export class WorkflowOrchestratorService {
       throw new Error(`Workflow with id "${workflowId}" not found.`)
     }
 
-    const flow = exportedWorkflow(container as MedusaContainer)
+    const flow = exportedWorkflow(
+      (container as MedusaContainer) ?? this.container_
+    )
 
     const events = this.buildWorkflowEvents({
       customEventHandlers: eventHandlers,

packages/modules/workflow-engine-redis/src/services/workflows-module.ts

@@ -60,6 +60,9 @@ export class WorkflowsModuleService<
     onApplicationPrepareShutdown: async () => {
       await this.workflowOrchestratorService_.onApplicationPrepareShutdown()
     },
+    onApplicationStart: async () => {
+      await this.workflowOrchestratorService_.onApplicationStart()
+    },
   }
 
   @InjectSharedContext()

packages/modules/workflow-engine-redis/src/utils/workflow-orchestrator-storage.ts

@@ -29,6 +29,8 @@ export class RedisDistributedTransactionStorage
   private workflowOrchestratorService_: WorkflowOrchestratorService
 
   private redisClient: Redis
+  private redisWorkerConnection: Redis
+  private queueName: string
   private queue: Queue
   private worker: Worker
 
@@ -47,12 +49,24 @@ export class RedisDistributedTransactionStorage
   }) {
     this.workflowExecutionService_ = workflowExecutionService
     this.logger_ = logger
-
     this.redisClient = redisConnection
-
+    this.redisWorkerConnection = redisWorkerConnection
+    this.queueName = redisQueueName
     this.queue = new Queue(redisQueueName, { connection: this.redisClient })
+  }
+
+  async onApplicationPrepareShutdown() {
+    // Close worker gracefully, i.e. wait for the current jobs to finish
+    await this.worker.close()
+  }
+
+  async onApplicationShutdown() {
+    await this.queue.close()
+  }
+
+  async onApplicationStart() {
     this.worker = new Worker(
-      redisQueueName,
+      this.queueName,
       async (job) => {
         const allJobs = [
           JobType.RETRY,
@@ -75,19 +89,10 @@ export class RedisDistributedTransactionStorage
           )
         }
       },
-      { connection: redisWorkerConnection }
+      { connection: this.redisWorkerConnection }
     )
   }
 
-  async onApplicationPrepareShutdown() {
-    // Close worker gracefully, i.e. wait for the current jobs to finish
-    await this.worker.close()
-  }
-
-  async onApplicationShutdown() {
-    await this.queue.close()
-  }
-
   setWorkflowOrchestratorService(workflowOrchestratorService) {
     this.workflowOrchestratorService_ = workflowOrchestratorService
   }